diff --git a/cecli/coders/base_coder.py b/cecli/coders/base_coder.py index 7b23d50a814..a5bde7ff756 100755 --- a/cecli/coders/base_coder.py +++ b/cecli/coders/base_coder.py @@ -2860,6 +2860,8 @@ async def format_in_executor(): should_retry = ex_info.retry if ex_info.name == "ServiceUnavailableError": should_retry = should_retry or retry_config["retry_on_unavailable"] + if ex_info.name == "PermissionDeniedError": + should_retry = should_retry or retry_config["retry_on_forbidden"] if should_retry: retry_delay *= retry_config["retry_backoff_factor"] diff --git a/cecli/models.py b/cecli/models.py index 619cf45e34c..6c87b7c7300 100644 --- a/cecli/models.py +++ b/cecli/models.py @@ -133,6 +133,7 @@ class ModelSettings: retries: Optional[dict] = None retry_backoff_factor: float = 1.5 retry_on_unavailable: bool = True + retry_on_forbidden: bool = False retry_timeout: float = 30 request_timeout: int = request_timeout debug: bool = False @@ -1454,6 +1455,7 @@ async def send_completion( retry_config = parse_retry_config(self.retries) self.retry_on_unavailable = retry_config["retry_on_unavailable"] + self.retry_on_forbidden = retry_config["retry_on_forbidden"] self.retry_backoff_factor = retry_config["retry_backoff_factor"] self.retry_timeout = retry_config["retry_timeout"] @@ -1486,6 +1488,8 @@ async def send_completion( should_retry = ex_info.retry if ex_info.name == "ServiceUnavailableError": should_retry = should_retry or self.retry_on_unavailable + elif ex_info.name == "PermissionDeniedError": + should_retry = should_retry or self.retry_on_forbidden custom_retry_delay = self._extract_retry_delay(err) if custom_retry_delay is not None: @@ -1550,6 +1554,7 @@ async def simple_send_with_retries( retry_backoff_factor = retry_config["retry_backoff_factor"] retry_timeout = retry_config["retry_timeout"] retry_on_unavailable = retry_config["retry_on_unavailable"] + retry_on_forbidden = retry_config["retry_on_forbidden"] if self.verbose: dump(messages) @@ -1603,6 +1608,8 @@ async def simple_send_with_retries( should_retry = ex_info.retry if ex_info.name == "ServiceUnavailableError": should_retry = should_retry or retry_on_unavailable + elif ex_info.name == "PermissionDeniedError": + should_retry = should_retry or retry_on_forbidden custom_retry_delay = self._extract_retry_delay(err) if custom_retry_delay is not None: @@ -1796,6 +1803,7 @@ def parse_retry_config(retries_input): retry_timeout: 30 retry_backoff_factor: 1.5 retry_on_unavailable: True + retry_on_forbidden: False retry_on_empty: False """ config = dict() @@ -1821,6 +1829,7 @@ def _get(key, default): "retry_timeout": float(_get("retry_timeout", 30)), "retry_backoff_factor": float(_get("retry_backoff_factor", 1.5)), "retry_on_unavailable": bool(_get("retry_on_unavailable", True)), + "retry_on_forbidden": bool(_get("retry_on_forbidden", False)), "retry_on_empty": bool(_get("retry_on_empty", False)), } diff --git a/cecli/sessions.py b/cecli/sessions.py new file mode 100644 index 00000000000..a8e178ec04e --- /dev/null +++ b/cecli/sessions.py @@ -0,0 +1,502 @@ +"""Session management utilities for cecli.""" + +import json +import os +from pathlib import Path +from typing import Dict, List, Optional + +from cecli import models +from cecli.decoding import safe_open +from cecli.helpers import crypto as session_crypto +from cecli.helpers.conversation import ConversationService, MessageTag + + +class SessionManager: + """Manages chat session saving, listing, and loading.""" + + def __init__(self, coder, io): + self.coder = coder + self.io = io + + def save_session(self, session_name: str, output=True) -> bool: + """Save the current chat session to a named file.""" + if not session_name: + if output: + self.io.tool_error("Please provide a session name.") + return False + + session_name = session_name.replace(".json", "") + session_dir = self._get_session_directory() + session_file = session_dir / f"{session_name}.json" + + if session_file.exists(): + if output: + self.io.tool_warning(f"Session '{session_name}' already exists. Overwriting.") + + try: + session_data = self._build_session_data(session_name) + if not self._write_session_file(session_file, session_data): + return False + + if output: + suffix = " (encrypted)" if self._session_encrypt_settings()[0] else "" + self.io.tool_output(f"Session saved: {session_file}{suffix}") + + return True + + except Exception as e: + self.io.tool_error(f"Error saving session: {e}") + return False + + def list_sessions(self) -> List[Dict]: + """List all saved sessions with metadata.""" + session_dir = self._get_session_directory() + session_files = list(session_dir.glob("*.json")) + + if not session_files: + self.io.tool_output("No saved sessions found.") + return [] + + sessions = [] + for session_file in sorted(session_files, key=lambda x: x.stat().st_mtime, reverse=True): + try: + raw = session_file.read_bytes() + if session_crypto.is_encrypted_payload(raw): + _, key = self._session_encrypt_settings() + if not key: + sessions.append( + { + "name": session_file.stem, + "file": session_file, + "model": "encrypted", + "edit_format": "—", + "num_messages": 0, + "num_files": 0, + "encrypted": True, + } + ) + continue + session_data = session_crypto.decrypt_session_bytes(raw, key) + else: + session_data = json.loads(raw.decode("utf-8")) + if not isinstance(session_data, dict): + raise ValueError("not a session object") + + session_info = { + "name": session_file.stem, + "file": session_file, + "model": session_data.get("model", "unknown"), + "edit_format": session_data.get("edit_format", "unknown"), + "num_messages": ( + len(session_data.get("chat_history", {}).get("done_messages", [])) + + len(session_data.get("chat_history", {}).get("cur_messages", [])) + ), + "num_files": ( + len(session_data.get("files", {}).get("editable", [])) + + len(session_data.get("files", {}).get("read_only", [])) + + len(session_data.get("files", {}).get("read_only_stubs", [])) + ), + "encrypted": session_crypto.is_encrypted_payload(raw), + } + sessions.append(session_info) + + except Exception as e: + self.io.tool_output(f" {session_file.stem} [error reading: {e}]") + + return sessions + + async def load_session(self, session_identifier: str, switch=True, quiet: bool = False) -> bool: + """Load a saved session by name or file path.""" + if not session_identifier: + self.io.tool_error("Please provide a session name or file path.") + return False + + # Try to find the session file + session_file = self._find_session_file(session_identifier) + if not session_file: + return False + + session_data = self._read_session_file(session_file, quiet=quiet) + if session_data is None: + return False + + if not isinstance(session_data, dict) or "version" not in session_data: + if not quiet: + self.io.tool_error("Invalid session format.") + return False + + # Apply session data + applied, loaded_edit_format = await self._apply_session_data(session_data, session_file) + if applied and switch: + from cecli.commands import SwitchCoderSignal + + edit_format_to_switch_to = self.coder.edit_format + if loaded_edit_format: + edit_format_to_switch_to = loaded_edit_format + self.coder.edit_format = loaded_edit_format + + raise SwitchCoderSignal( + edit_format=edit_format_to_switch_to, + from_coder=self.coder, + summarize_from_coder=False, + show_announcements=True, + ) + return applied + + def _get_session_directory(self) -> Path: + """Get the session directory, creating it if necessary.""" + session_dir = Path(self.coder.abs_root_path(".cecli/sessions")) + os.makedirs(session_dir, exist_ok=True) + return session_dir + + def _session_encrypt_settings(self) -> tuple[bool, bytes | None]: + args = getattr(self.coder, "args", None) + if not args or not getattr(args, "session_encrypt", False): + return False, None + key_file = getattr(args, "session_key_file", None) + return True, session_crypto.resolve_key(key_file=key_file) + + def _read_session_file(self, session_file: Path, quiet: bool = False) -> dict | None: + try: + data = session_file.read_bytes() + except OSError as e: + if not quiet: + self.io.tool_error(f"Error reading session: {e}") + return None + try: + if session_crypto.is_encrypted_payload(data): + args = getattr(self.coder, "args", None) + key_file = getattr(args, "session_key_file", None) if args else None + key = session_crypto.resolve_key(key_file=key_file) + if not key: + if not quiet: + self.io.tool_error( + "Session is encrypted but no key is configured " + f"({session_crypto.KEY_ENV} or --session-key-file)." + ) + return None + return session_crypto.decrypt_session_bytes(data, key) + parsed = json.loads(data.decode("utf-8")) + if not isinstance(parsed, dict): + if not quiet: + self.io.tool_error("Invalid session format.") + return None + return parsed + except session_crypto.SessionCryptoError as e: + if not quiet: + self.io.tool_error(str(e)) + return None + except (UnicodeDecodeError, json.JSONDecodeError) as e: + if not quiet: + self.io.tool_error(f"Error loading session: {e}") + return None + + def _write_session_file(self, session_file: Path, session_data: dict) -> bool: + encrypt_enabled, key = self._session_encrypt_settings() + try: + if encrypt_enabled: + if not key: + self.io.tool_error( + "Session encryption is enabled but no key is configured " + f"({session_crypto.KEY_ENV} or --session-key-file)." + ) + return False + session_file.write_bytes(session_crypto.encrypt_session_dict(session_data, key)) + else: + with safe_open(session_file, "w") as f: + json.dump(session_data, f, indent=2) + return True + except session_crypto.SessionCryptoError as e: + self.io.tool_error(str(e)) + return False + except OSError as e: + self.io.tool_error(f"Error saving session: {e}") + return False + + def _build_session_data(self, session_name) -> Dict: + """Build session data dictionary from current coder state.""" + # Get relative paths for all files + editable_files = [ + self.coder.get_rel_fname(abs_fname) for abs_fname in self.coder.abs_fnames + ] + read_only_files = [ + self.coder.get_rel_fname(abs_fname) for abs_fname in self.coder.abs_read_only_fnames + ] + read_only_stubs_files = [ + self.coder.get_rel_fname(abs_fname) + for abs_fname in self.coder.abs_read_only_stubs_fnames + ] + + # Capture todo list content so it can be restored with the session + todo_content = None + try: + todo_path = self.coder.abs_root_path(self.coder.local_agent_folder("todo.txt")) + if os.path.isfile(todo_path): + todo_content = self.io.read_text(todo_path) + if todo_content is None: + todo_content = "" + except Exception as e: + self.io.tool_warning(f"Could not read todo list file: {e}") + + # Get CUR and DONE messages from ConversationManager + connected_mcps = [] + if hasattr(self.coder, "mcp_manager") and self.coder.mcp_manager: + connected_mcps = [server.name for server in self.coder.mcp_manager.connected_servers] + + # Get CUR and DONE messages from ConversationManager + connected_mcps = [] + if hasattr(self.coder, "mcp_manager") and self.coder.mcp_manager: + connected_mcps = [server.name for server in self.coder.mcp_manager.connected_servers] + + skills_data = None + if hasattr(self.coder, "skills_manager") and self.coder.skills_manager: + skills_data = { + "skills_paths": [str(p) for p in self.coder.skills_manager.directory_paths], + "skills_includelist": ( + list(self.coder.skills_manager.include_list) + if self.coder.skills_manager.include_list is not None + else [] + ), + "skills_excludelist": ( + list(self.coder.skills_manager.exclude_list) + if self.coder.skills_manager.exclude_list is not None + else [] + ), + } + + agent_config_data = None + if hasattr(self.coder, "agent_config"): + agent_config_data = { + "tools_paths": self.coder.agent_config.get("tools_paths", []), + "tools_includelist": self.coder.agent_config.get("tools_includelist", []), + "tools_excludelist": self.coder.agent_config.get("tools_excludelist", []), + } + + # Flush any queued messages so the saved chat history is complete + ConversationService.get_manager(self.coder).flush_queue() + + return { + "version": 1, + "session_name": session_name, + "model": self.coder.main_model.name, + "weak_model": self.coder.main_model.weak_model.name, + "editor_model": self.coder.main_model.editor_model.name, + "agent_model": self.coder.main_model.agent_model.name, + "editor_edit_format": self.coder.main_model.editor_edit_format, + "edit_format": self.coder.edit_format, + "chat_history": { + "done_messages": ( + ConversationService.get_manager(self.coder).get_messages_dict(MessageTag.DONE) + ), + "cur_messages": ( + ConversationService.get_manager(self.coder).get_messages_dict(MessageTag.CUR) + ), + }, + "files": { + "editable": editable_files, + "read_only": read_only_files, + "read_only_stubs": read_only_stubs_files, + }, + "settings": { + "auto_commits": self.coder.auto_commits, + "auto_lint": self.coder.auto_lint, + "auto_test": self.coder.auto_test, + }, + "todo_list": todo_content, + "mcps": connected_mcps, + "skills": skills_data, + "tools": agent_config_data, + "usage": { + "total_tokens_sent": self.coder.total_tokens_sent, + "total_tokens_received": self.coder.total_tokens_received, + "total_cached_tokens": self.coder.total_cached_tokens, + "total_cost": self.coder.total_cost, + }, + } + + def _find_session_file(self, session_identifier: str) -> Optional[Path]: + """Find session file by name or path.""" + # Check if it's a direct file path + session_file = Path(session_identifier) + if session_file.exists(): + return session_file + + # Check if it's a session name in the sessions directory + session_dir = self._get_session_directory() + + # Try with .json extension + if not session_identifier.endswith(".json"): + session_file = session_dir / f"{session_identifier}.json" + if session_file.exists(): + return session_file + + session_file = session_dir / f"{session_identifier}" + if session_file.exists(): + return session_file + + self.io.tool_error(f"Session not found: {session_identifier}") + self.io.tool_output("Use /list-sessions to see available sessions.") + return None + + async def _apply_session_data( + self, session_data: Dict, session_file: Path + ) -> (bool, Optional[str]): + """Apply session data to current coder state. + + Returns: + A tuple of (success, edit_format) + """ + try: + # Clear current state + self.coder.abs_fnames = set() + self.coder.abs_read_only_fnames = set() + self.coder.abs_read_only_stubs_fnames = set() + + # Load files + files = session_data.get("files", {}) + for rel_fname in files.get("editable", []): + abs_fname = self.coder.abs_root_path(rel_fname) + if os.path.exists(abs_fname): + self.coder.abs_fnames.add(abs_fname) + else: + self.io.tool_warning(f"File not found, skipping: {rel_fname}") + + for rel_fname in files.get("read_only", []): + abs_fname = self.coder.abs_root_path(rel_fname) + if os.path.exists(abs_fname): + self.coder.abs_read_only_fnames.add(abs_fname) + else: + self.io.tool_warning(f"File not found, skipping: {rel_fname}") + + for rel_fname in files.get("read_only_stubs", []): + abs_fname = self.coder.abs_root_path(rel_fname) + if os.path.exists(abs_fname): + self.coder.abs_read_only_stubs_fnames.add(abs_fname) + else: + self.io.tool_warning(f"File not found, skipping: {rel_fname}") + + # Load usage stats + usage = session_data.get("usage", {}) + self.coder.total_tokens_sent = usage.get("total_tokens_sent", 0) + self.coder.total_tokens_received = usage.get("total_tokens_received", 0) + self.coder.total_cached_tokens = usage.get("total_cached_tokens", 0) + self.coder.total_cost = usage.get("total_cost", 0.0) + if session_data.get("model"): + self.coder.main_model = models.Model( + session_data.get("model", self.coder.args.model), + weak_model=session_data.get("weak_model", self.coder.args.weak_model), + editor_model=session_data.get("editor_model", self.coder.args.editor_model), + agent_model=session_data.get("agent_model", self.coder.args.agent_model), + editor_edit_format=session_data.get( + "editor_edit_format", self.coder.args.editor_edit_format + ), + io=self.io, + verbose=self.coder.args.verbose, + retries=self.coder.main_model.retries, + debug=self.coder.main_model.debug, + ) + + # Load settings + settings = session_data.get("settings", {}) + if "auto_commits" in settings: + self.coder.auto_commits = settings["auto_commits"] + if "auto_lint" in settings: + self.coder.auto_lint = settings["auto_lint"] + if "auto_test" in settings: + self.coder.auto_test = settings["auto_test"] + + # Restore todo list content if present in the session + if "todo_list" in session_data: + todo_path = self.coder.abs_root_path(self.coder.local_agent_folder("todo.txt")) + todo_content = session_data.get("todo_list") + try: + if todo_content is None: + if os.path.exists(todo_path): + os.remove(todo_path) + else: + self.io.write_text(todo_path, todo_content) + except Exception as e: + self.io.tool_warning(f"Could not restore todo list: {e}") + + # Clear CUR and DONE messages from ConversationManager + ConversationService.get_manager(self.coder).reset() + ConversationService.get_files(self.coder).reset() + self.coder.format_chat_chunks() + + # Load chat history + chat_history = session_data.get("chat_history", {}) + done_messages = chat_history.get("done_messages", []) + cur_messages = chat_history.get("cur_messages", []) + + # Add messages to ConversationManager (source of truth) + # Add done messages + for msg in done_messages: + ConversationService.get_manager(self.coder).add_message( + message_dict=msg, + tag=MessageTag.DONE, + ) + # Add current messages + for msg in cur_messages: + ConversationService.get_manager(self.coder).add_message( + message_dict=msg, + tag=MessageTag.CUR, + ) + + self.io.tool_output( + f"Session loaded: {session_data.get('session_name', session_file.stem)}" + ) + self.io.tool_output( + f"Model: {session_data.get('model', 'unknown')}, Edit format:" + f" {session_data.get('edit_format', 'unknown')}" + ) + + # Show summary + num_messages = len(self.coder.done_messages) + len(self.coder.cur_messages) + num_files = ( + len(self.coder.abs_fnames) + + len(self.coder.abs_read_only_fnames) + + len(self.coder.abs_read_only_stubs_fnames) + ) + self.io.tool_output(f"Loaded {num_messages} messages and {num_files} files") + + # Load MCPs + saved_mcps = session_data.get("mcps", []) + if hasattr(self.coder, "mcp_manager") and self.coder.mcp_manager: + current_mcps = {server.name for server in self.coder.mcp_manager.connected_servers} + saved_mcps_set = set(saved_mcps) + + to_disconnect = current_mcps - saved_mcps_set + for mcp_name in to_disconnect: + await self.coder.mcp_manager.disconnect_server(mcp_name) + + to_connect = saved_mcps_set - current_mcps + for mcp_name in to_connect: + await self.coder.mcp_manager.connect_server(mcp_name) + + # Load skills + skills_data = session_data.get("skills") + if skills_data and hasattr(self.coder, "skills_manager") and self.coder.skills_manager: + self.coder.skills_manager.directory_paths = skills_data.get("skills_paths", []) + self.coder.skills_manager.include_list = set( + skills_data.get("skills_includelist", []) + ) + self.coder.skills_manager.exclude_list = set( + skills_data.get("skills_excludelist", []) + ) + + # Load tools config + agent_config_data = session_data.get("tools") + if agent_config_data and hasattr(self.coder, "agent_config"): + self.coder.agent_config.update(agent_config_data) + from cecli.tools.utils.registry import ToolRegistry + + ToolRegistry.build_registry(agent_config=self.coder.agent_config) + self.coder.loaded_custom_tools = ToolRegistry.loaded_custom_tools + + # Return True and the edit format so the Coder can be switched + edit_format = session_data.get("edit_format") + return True, edit_format + + except Exception as e: + self.io.tool_error(f"Error applying session data: {e}") + return False, None diff --git a/requirements.txt b/requirements.txt index 64d185820fb..9e182322f7d 100644 --- a/requirements.txt +++ b/requirements.txt @@ -608,7 +608,7 @@ uvicorn[standard]==0.38.0 # -c requirements/common-constraints.txt # chromadb # mcp -uvloop==0.22.1 +uvloop==0.22.1 ; platform_python_implementation != 'PyPy' and sys_platform != 'cygwin' and sys_platform != 'win32' # via # -c requirements/common-constraints.txt # uvicorn @@ -642,6 +642,6 @@ zipp==3.23.0 # via # -c requirements/common-constraints.txt # importlib-metadata - + tree-sitter==0.23.2; python_version < "3.10" tree-sitter>=0.25.1; python_version >= "3.10" diff --git a/requirements/common-constraints.txt b/requirements/common-constraints.txt index 5f79887e8e4..53f2548b1f7 100644 --- a/requirements/common-constraints.txt +++ b/requirements/common-constraints.txt @@ -514,7 +514,7 @@ uvicorn[standard]==0.38.0 # via # chromadb # mcp -uvloop==0.22.1 +uvloop==0.22.1 ; platform_python_implementation != 'PyPy' and sys_platform != 'cygwin' and sys_platform != 'win32' # via uvicorn virtualenv==20.35.4 # via pre-commit diff --git a/tests/basic/test_retry_config.py b/tests/basic/test_retry_config.py new file mode 100644 index 00000000000..56b609ea209 --- /dev/null +++ b/tests/basic/test_retry_config.py @@ -0,0 +1,68 @@ +from unittest.mock import AsyncMock, call, patch + +import pytest + +from cecli.llm import litellm +from cecli.models import Model, _parse_retry_config + + +def test_parse_retry_config_string(): + config_str = '{"retry_timeout": 15, "retry-on-empty": true}' + result = _parse_retry_config(config_str) + assert result["retry_timeout"] == 15.0 + assert result["retry_on_empty"] is True + # defaults + assert result["retry_backoff_factor"] == 1.5 + assert result["retry_on_unavailable"] is True + + +def test_parse_retry_config_dict(): + config_dict = { + "retry_timeout": 10.0, + "retry_backoff_factor": 2.0, + "retry-on-unavailable": False, + } + result = _parse_retry_config(config_dict) + assert result["retry_timeout"] == 10.0 + assert result["retry_backoff_factor"] == 2.0 + assert result["retry_on_unavailable"] is False + assert result["retry_on_empty"] is False + + +@pytest.mark.asyncio +async def test_simple_send_with_retries_honors_timeout(): + # Setup model with a short retry timeout limit + model = Model("gpt-4o", retries={"retry_timeout": 0.5, "retry_backoff_factor": 2.0}) + + # retry_delay starts at 0.125 and is multiplied by the backoff factor + # BEFORE each retry sleep; retry_timeout caps the per-retry delay: + # attempt 1 fails -> 0.125 * 2.0 = 0.25 (<= 0.5, sleep and retry) + # attempt 2 fails -> 0.25 * 2.0 = 0.50 (<= 0.5, sleep and retry) + # attempt 3 fails -> 0.50 * 2.0 = 1.00 (> 0.5, give up) + err = litellm.APIConnectionError( + message="Simulated connection error", + llm_provider="openai", + model="gpt-4o", + request=None, + ) + + mock_send = AsyncMock(side_effect=err) + + with ( + patch.object(model, "send_completion", mock_send), + patch("time.sleep") as mock_sleep, + patch("builtins.print"), + ): # Mute prints in test output + + content, response = await model.simple_send_with_retries(messages=[]) + + # It should exit yielding None, None because it exhausted retries. + assert content is None + assert response is None + + # The backoff factor is applied before each sleep, so the sleeps are + # 0.25 then 0.50; the third failure would need 1.0 > 0.5, so it stops. + + assert mock_send.call_count == 3 + assert mock_sleep.call_count == 2 + assert mock_sleep.call_args_list == [call(0.25), call(0.5)]