code improvements for styling and validation

Author: Murat Kaan Meral

src/strands_agents_builder/strands.py

@@ -16,6 +16,7 @@
 from strands_agents_builder.utils import model_utils
 from strands_agents_builder.utils.kb_utils import load_system_prompt, store_conversation_in_kb
 from strands_agents_builder.utils.session_utils import (
+    console,
     display_agent_history,
     handle_session_commands,
     list_sessions_command,
@@ -50,7 +51,7 @@ def handle_shell_command(agent: Agent, command: str, user_input: str) -> None:
             non_interactive_mode=True,
         )
     except Exception as e:
-        print(f"Error: {str(e)}")
+        console.print(f"[red]Error: {str(e)}[/red]")
 
 
 def execute_interactive_mode(
@@ -120,7 +121,7 @@ def execute_interactive_mode(
             break
         except Exception as e:
             callback_handler(force_stop=True)  # Stop spinners
-            print(f"Error: {str(e)}")
+            console.print(f"[red]Error: {str(e)}[/red]")
 
 
 def main():

src/strands_agents_builder/utils/session_utils.py

@@ -4,6 +4,7 @@
 """
 
 import datetime
+import logging
 import time
 import uuid
 from pathlib import Path
@@ -19,6 +20,42 @@
 # Create console for rich formatting
 console = Console()
 
+# Constants
+SESSION_PREFIX = "session_"
+DEFAULT_DISPLAY_LIMIT = 10
+
+# Set up logging
+logger = logging.getLogger(__name__)
+
+
+def validate_session_id(session_id: str) -> bool:
+    """Validate that a session ID is safe to use as a directory name."""
+    if not session_id:
+        return False
+
+    # Check for basic safety - no path separators, no hidden files, reasonable length
+    if any(char in session_id for char in ["/", "\\", "..", "\0"]):
+        return False
+
+    if session_id.startswith(".") or len(session_id) > 255:
+        return False
+
+    return True
+
+
+def validate_session_path(path: str) -> bool:
+    """Validate that a session path is safe to use."""
+    if not path:
+        return False
+
+    try:
+        # Try to create a Path object and check if it's absolute or relative
+        path_obj = Path(path)
+        # Basic validation - path should be reasonable
+        return len(str(path_obj)) < 4096  # Reasonable path length limit
+    except (ValueError, OSError):
+        return False
+
 
 def generate_session_id() -> str:
     """Generate a unique session ID based on timestamp and UUID."""
@@ -44,11 +81,14 @@ def create_session_manager(
     session_id: Optional[str] = None, base_path: Optional[str] = None
 ) -> Optional[FileSessionManager]:
     """Create a FileSessionManager with the given or generated session ID. Returns None if no base_path."""
-    if not base_path:
+    if not base_path or not validate_session_path(base_path):
         return None
 
     if session_id is None:
         session_id = generate_session_id()
+    elif not validate_session_id(session_id):
+        logger.warning(f"Invalid session ID provided: {session_id}")
+        return None
 
     # Create the sessions directory since we're actually creating a session manager
     sessions_dir = get_sessions_directory(base_path, create=True)
@@ -57,68 +97,81 @@ def create_session_manager(
 
 def list_available_sessions(base_path: Optional[str] = None) -> list[str]:
     """List all available session IDs in the sessions directory."""
-    if not base_path:
+    if not base_path or not validate_session_path(base_path):
         return []
 
     # Don't create directory, just check if it exists
     sessions_dir = get_sessions_directory(base_path, create=False)
     session_ids = []
 
     if sessions_dir and sessions_dir.exists():
-        for session_dir in sessions_dir.iterdir():
-            if session_dir.is_dir() and session_dir.name.startswith("session_"):
-                # Extract session ID from directory name (remove "session_" prefix)
-                session_id = session_dir.name[8:]  # len("session_") = 8
-                session_ids.append(session_id)
+        try:
+            for session_dir in sessions_dir.iterdir():
+                if session_dir.is_dir() and session_dir.name.startswith(SESSION_PREFIX):
+                    # Extract session ID from directory name (remove "session_" prefix)
+                    session_id = session_dir.name[len(SESSION_PREFIX) :]
+                    if validate_session_id(session_id):
+                        session_ids.append(session_id)
+        except (OSError, PermissionError) as e:
+            logger.warning(f"Failed to list sessions in {base_path}: {e}")
 
     return sorted(session_ids)
 
 
 def session_exists(session_id: str, base_path: Optional[str] = None) -> bool:
     """Check if a session exists."""
-    if not base_path:
+    if not base_path or not validate_session_path(base_path) or not validate_session_id(session_id):
         return False
 
     # Don't create directory, just check if session exists
     sessions_dir = get_sessions_directory(base_path, create=False)
     if not sessions_dir:
         return False
 
-    session_dir = sessions_dir / f"session_{session_id}"
+    session_dir = sessions_dir / f"{SESSION_PREFIX}{session_id}"
     return session_dir.exists() and (session_dir / "session.json").exists()
 
 
 def get_session_info(session_id: str, base_path: Optional[str] = None) -> Optional[dict]:
     """Get basic information about a session."""
-    if not base_path or not session_exists(session_id, base_path):
+    if not base_path or not validate_session_path(base_path) or not validate_session_id(session_id):
+        return None
+
+    if not session_exists(session_id, base_path):
         return None
 
     # Don't create directory, just get the path
     sessions_dir = get_sessions_directory(base_path, create=False)
     if not sessions_dir:
         return None
 
-    session_dir = sessions_dir / f"session_{session_id}"
-
-    # Get creation time from directory
-    created_at = session_dir.stat().st_ctime
+    session_dir = sessions_dir / f"{SESSION_PREFIX}{session_id}"
 
-    # Count messages across all agents
-    total_messages = 0
-    agents_dir = session_dir / "agents"
-    if agents_dir.exists():
-        for agent_dir in agents_dir.iterdir():
-            if agent_dir.is_dir():
-                messages_dir = agent_dir / "messages"
-                if messages_dir.exists():
-                    total_messages += len([f for f in messages_dir.iterdir() if f.is_file() and f.suffix == ".json"])
-
-    return {
-        "session_id": session_id,
-        "created_at": created_at,
-        "total_messages": total_messages,
-        "path": str(session_dir),
-    }
+    try:
+        # Get creation time from directory
+        created_at = session_dir.stat().st_ctime
+
+        # Count messages across all agents
+        total_messages = 0
+        agents_dir = session_dir / "agents"
+        if agents_dir.exists():
+            for agent_dir in agents_dir.iterdir():
+                if agent_dir.is_dir():
+                    messages_dir = agent_dir / "messages"
+                    if messages_dir.exists():
+                        total_messages += len(
+                            [f for f in messages_dir.iterdir() if f.is_file() and f.suffix == ".json"]
+                        )
+
+        return {
+            "session_id": session_id,
+            "created_at": created_at,
+            "total_messages": total_messages,
+            "path": str(session_dir),
+        }
+    except (OSError, PermissionError) as e:
+        logger.warning(f"Failed to get session info for {session_id}: {e}")
+        return None
 
 
 def list_sessions_command(session_base_path: Optional[str]) -> None:
@@ -146,8 +199,8 @@ def display_agent_history(agent, session_id: str) -> None:
     """Display conversation history from an agent's loaded messages."""
     try:
         if agent.messages and len(agent.messages) > 0:
-            # Display last 10 messages (5 pairs) completely
-            display_limit = 10
+            # Display last messages completely
+            display_limit = DEFAULT_DISPLAY_LIMIT
 
             # Create header message
             header_text = f"Resuming session: {session_id}"
@@ -192,14 +245,15 @@ def display_agent_history(agent, session_id: str) -> None:
                     print(f"{Fore.WHITE}{content}{Style.RESET_ALL}")
                     print()  # Empty line after assistant message
 
-    except Exception:
-        # If we can't load history, just continue silently
-        pass
+    except Exception as e:
+        # If we can't load history, log the error but continue
+        logger.warning(f"Failed to display agent history for session {session_id}: {e}")
+        console.print("[yellow]Warning: Could not load session history[/yellow]")
 
 
 def setup_session_management(
     session_id: Optional[str], session_base_path: Optional[str]
-) -> Tuple[Optional[object], Optional[str], bool]:
+) -> Tuple[Optional[FileSessionManager], Optional[str], bool]:
     """Set up session management if enabled. Returns (session_manager, session_id, is_resuming)."""
     session_manager = None
     resolved_session_id = None
@@ -228,32 +282,34 @@ def handle_session_commands(command: str, session_id: Optional[str], session_bas
         info = get_session_info(session_id, session_base_path)
         if info:
             created = datetime.datetime.fromtimestamp(info["created_at"]).strftime("%Y-%m-%d %H:%M:%S")
-            print(f"Session ID: {info['session_id']}")
-            print(f"Created: {created}")
-            print(f"Total messages: {info['total_messages']}")
+            console.print(f"[bold cyan]Session ID:[/bold cyan] {info['session_id']}")
+            console.print(f"[bold cyan]Created:[/bold cyan] {created}")
+            console.print(f"[bold cyan]Total messages:[/bold cyan] {info['total_messages']}")
         return True
 
     elif command == "session list" and session_base_path:
         sessions = list_available_sessions(session_base_path)
         if not sessions:
-            print("No sessions found.")
+            console.print("[yellow]No sessions found.[/yellow]")
         else:
-            print("Available sessions:")
+            console.print("[bold cyan]Available sessions:[/bold cyan]")
             for sid in sessions:
                 info = get_session_info(sid, session_base_path)
                 if info:
                     created = datetime.datetime.fromtimestamp(info["created_at"]).strftime("%Y-%m-%d %H:%M:%S")
-                    current = " (current)" if sid == session_id else ""
-                    print(f"  {sid} (created: {created}, messages: {info['total_messages']}){current}")
+                    current = " [dim](current)[/dim]" if sid == session_id else ""
+                    console.print(
+                        f"  [green]{sid}[/green] (created: {created}, messages: {info['total_messages']}){current}"
+                    )
         return True
 
     elif command.startswith("session "):
         if session_base_path:
-            print("Available session commands:")
-            print("  !session info  - Show current session details")
-            print("  !session list  - List all available sessions")
+            console.print("[bold cyan]Available session commands:[/bold cyan]")
+            console.print("  [green]!session info[/green]  - Show current session details")
+            console.print("  [green]!session list[/green]  - List all available sessions")
         else:
-            print("Error: Session management not enabled.")
+            console.print("[red]Error: Session management not enabled.[/red]")
         return True
 
     return False

tests/test_strands.py

@@ -172,11 +172,11 @@ def test_eof_error_exception(self, mock_goodbye, mock_agent, mock_input):
         # Verify goodbye message was called
         mock_goodbye.assert_called_once()
 
-    @mock.patch("builtins.print")
+    @mock.patch("strands_agents_builder.utils.session_utils.console.print")
     @mock.patch.object(strands, "get_user_input")
     @mock.patch.object(strands, "Agent")
     @mock.patch.object(strands, "callback_handler")
-    def test_general_exception_handling(self, mock_callback_handler, mock_agent, mock_input, mock_print):
+    def test_general_exception_handling(self, mock_callback_handler, mock_agent, mock_input, mock_console_print):
         """Test handling of general exceptions in interactive mode"""
         # Setup mocks
         mock_agent_instance = mock.MagicMock()
@@ -194,7 +194,7 @@ def test_general_exception_handling(self, mock_callback_handler, mock_agent, moc
             strands.main()
 
         # Verify error was called
-        mock_print.assert_any_call("Error: Test error")
+        mock_console_print.assert_any_call("[red]Error: Test error[/red]")
 
         # Verify callback_handler was called to stop spinners
         mock_callback_handler.assert_called_once_with(force_stop=True)
@@ -314,9 +314,16 @@ def test_general_exception(self, mock_agent, mock_bedrock, mock_load_prompt, mon
 class TestShellCommandError:
     """Test shell command error handling"""
 
-    @mock.patch("builtins.print")
+    @mock.patch("strands_agents_builder.utils.session_utils.console.print")
     def test_shell_command_exception(
-        self, mock_print, mock_agent, mock_bedrock, mock_load_prompt, mock_user_input, mock_welcome_message, monkeypatch
+        self,
+        mock_console_print,
+        mock_agent,
+        mock_bedrock,
+        mock_load_prompt,
+        mock_user_input,
+        mock_welcome_message,
+        monkeypatch,
     ):
         """Test handling exceptions when executing shell commands"""
         # Setup mocks
@@ -334,7 +341,7 @@ def test_shell_command_exception(
             strands.main()
 
         # Verify error was called
-        mock_print.assert_any_call("Error: Shell command failed")
+        mock_console_print.assert_any_call("[red]Error: Shell command failed[/red]")
 
 
 class TestKnowledgeBaseIntegration:
@@ -573,14 +580,14 @@ def test_agent_creation_without_session_manager(self, mock_create_manager, mock_
             call_kwargs = mock_agent_class.call_args[1]
             assert "session_manager" not in call_kwargs or call_kwargs.get("session_manager") is None
 
-    @mock.patch("builtins.print")
+    @mock.patch("strands_agents_builder.utils.session_utils.console.print")
     @mock.patch("strands_agents_builder.utils.session_utils.create_session_manager")
     @mock.patch("strands_agents_builder.utils.session_utils.get_session_info")
     def test_session_commands_in_interactive_mode(
         self,
         mock_get_info,
         mock_create_manager,
-        mock_print,
+        mock_console_print,
         mock_agent,
         mock_bedrock,
         mock_load_prompt,