diff --git a/.vscode/launch.json b/.vscode/launch.json index 661c24e..2993318 100644 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -1,5 +1,5 @@ { - "version": "2.0.1", + "version": "2.0.2", "configurations": [ { "name": "ACP Server", diff --git a/CHANGELOG.md b/CHANGELOG.md index fb33e59..363d32b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,11 +5,32 @@ All notable changes to this project will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [2.0.2] - 2026-01-30 + +### Added + +- Allow environment variables to be overridden by dotenv files +- Display custom rate limit messages depending on plan type + +### Changed + +- Made plan offer message more discreet in UI +- Speed up latest session scan and harden validation +- Updated pytest-xdist configuration to schedule single test chunks + +### Fixed + +- Prevent duplicate messages in persisted sessions +- Fix ACP bash tool to pass full command string for chained commands +- Fix global agent prompt not being loaded correctly +- Do not propose to "resume" when there is nothing to resume + + ## [2.0.1] - 2026-01-28 ### Fixed -- Encoding issues in Windows +- Fix encoding issues in Windows ## [2.0.0] - 2026-01-27 diff --git a/distribution/zed/extension.toml b/distribution/zed/extension.toml index b898d3a..473a2e5 100644 --- a/distribution/zed/extension.toml +++ b/distribution/zed/extension.toml @@ -1,7 +1,7 @@ id = "mistral-vibe" name = "Mistral Vibe" description = "Mistral's open-source coding assistant" -version = "2.0.1" +version = "2.0.2" schema_version = 1 authors = ["Mistral AI"] repository = "https://github.com/mistralai/mistral-vibe" @@ -11,25 +11,25 @@ name = "Mistral Vibe" icon = "./icons/mistral_vibe.svg" [agent_servers.mistral-vibe.targets.darwin-aarch64] -archive = "https://github.com/mistralai/mistral-vibe/releases/download/v2.0.1/vibe-acp-darwin-aarch64-2.0.1.zip" +archive = "https://github.com/mistralai/mistral-vibe/releases/download/v2.0.2/vibe-acp-darwin-aarch64-2.0.2.zip" cmd = "./vibe-acp" [agent_servers.mistral-vibe.targets.darwin-x86_64] -archive = "https://github.com/mistralai/mistral-vibe/releases/download/v2.0.1/vibe-acp-darwin-x86_64-2.0.1.zip" +archive = "https://github.com/mistralai/mistral-vibe/releases/download/v2.0.2/vibe-acp-darwin-x86_64-2.0.2.zip" cmd = "./vibe-acp" [agent_servers.mistral-vibe.targets.linux-aarch64] -archive = "https://github.com/mistralai/mistral-vibe/releases/download/v2.0.1/vibe-acp-linux-aarch64-2.0.1.zip" +archive = "https://github.com/mistralai/mistral-vibe/releases/download/v2.0.2/vibe-acp-linux-aarch64-2.0.2.zip" cmd = "./vibe-acp" [agent_servers.mistral-vibe.targets.linux-x86_64] -archive = "https://github.com/mistralai/mistral-vibe/releases/download/v2.0.1/vibe-acp-linux-x86_64-2.0.1.zip" +archive = "https://github.com/mistralai/mistral-vibe/releases/download/v2.0.2/vibe-acp-linux-x86_64-2.0.2.zip" cmd = "./vibe-acp" [agent_servers.mistral-vibe.targets.windows-aarch64] -archive = "https://github.com/mistralai/mistral-vibe/releases/download/v2.0.1/vibe-acp-windows-aarch64-2.0.1.zip" +archive = "https://github.com/mistralai/mistral-vibe/releases/download/v2.0.2/vibe-acp-windows-aarch64-2.0.2.zip" cmd = "./vibe-acp.exe" [agent_servers.mistral-vibe.targets.windows-x86_64] -archive = "https://github.com/mistralai/mistral-vibe/releases/download/v2.0.1/vibe-acp-windows-x86_64-2.0.1.zip" +archive = "https://github.com/mistralai/mistral-vibe/releases/download/v2.0.2/vibe-acp-windows-x86_64-2.0.2.zip" cmd = "./vibe-acp.exe" diff --git a/pyproject.toml b/pyproject.toml index 4616054..043ce92 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "mistral-vibe" -version = "2.0.1" +version = "2.0.2" description = "Minimal CLI coding agent by Mistral" readme = "README.md" requires-python = ">=3.12" @@ -159,5 +159,5 @@ max-nested-blocks = 4 ignore_decorators = ["@*"] [tool.pytest.ini_options] -addopts = "-vvvv -q -n auto --durations=5 --import-mode=importlib" +addopts = "-vvvv -q -n auto --durations=10 --import-mode=importlib --maxschedchunk=1" timeout = 10 diff --git a/tests/acp/test_bash.py b/tests/acp/test_bash.py index cb10c1d..3b3de6e 100644 --- a/tests/acp/test_bash.py +++ b/tests/acp/test_bash.py @@ -125,42 +125,6 @@ class TestAcpBashBasic: display = Bash.get_summary(args) assert display == "ls (timeout 10s)" - def test_parse_command_simple(self) -> None: - tool = Bash(config=BashToolConfig(), state=AcpBashState()) - env, command, args = tool._parse_command("ls") - assert env == [] - assert command == "ls" - assert args == [] - - def test_parse_command_with_args(self) -> None: - tool = Bash(config=BashToolConfig(), state=AcpBashState()) - env, command, args = tool._parse_command("ls -la src") - assert env == [] - assert command == "ls" - assert args == ["-la", "src"] - - def test_parse_command_with_env(self) -> None: - tool = Bash(config=BashToolConfig(), state=AcpBashState()) - env, command, args = tool._parse_command("NODE_ENV=test DEBUG=1 npm test") - assert len(env) == 2 - assert env[0].name == "NODE_ENV" - assert env[0].value == "test" - assert env[1].name == "DEBUG" - assert env[1].value == "1" - assert command == "npm" - assert args == ["test"] - - def test_parse_command_with_env_value_contains_equals(self) -> None: - tool = Bash(config=BashToolConfig(), state=AcpBashState()) - env, command, args = tool._parse_command( - "PATH=/usr/bin:/usr/local/bin echo hello" - ) - assert len(env) == 1 - assert env[0].name == "PATH" - assert env[0].value == "/usr/bin:/usr/local/bin" - assert command == "echo" - assert args == ["hello"] - class TestAcpBashExecution: @pytest.mark.asyncio @@ -181,8 +145,7 @@ class TestAcpBashExecution: # Verify create_terminal was called correctly params = mock_client._last_create_params assert params["session_id"] == "test_session_123" - assert params["command"] == "echo" - assert params["args"] == ["hello"] + assert params["command"] == "echo hello" assert params["cwd"] == str(Path.cwd()) # effective_workdir defaults to cwd @pytest.mark.asyncio @@ -200,16 +163,7 @@ class TestAcpBashExecution: await collect_result(tool.run(args)) params = mock_client._last_create_params - env = params["env"] - assert env is not None - assert ( - isinstance(env, list) and len(env) > 0 and isinstance(env[0], EnvVariable) - ) - assert len(env) == 1 - assert env[0].name == "NODE_ENV" - assert env[0].value == "test" - assert params["command"] == "npm" - assert params["args"] == ["run", "build"] + assert params["command"] == "NODE_ENV=test npm run build" @pytest.mark.asyncio async def test_run_with_nonzero_exit_code(self, mock_client: MockClient) -> None: diff --git a/tests/acp/test_initialize.py b/tests/acp/test_initialize.py index df9a344..6900f11 100644 --- a/tests/acp/test_initialize.py +++ b/tests/acp/test_initialize.py @@ -25,7 +25,7 @@ class TestACPInitialize: ), ) assert response.agent_info == Implementation( - name="@mistralai/mistral-vibe", title="Mistral Vibe", version="2.0.1" + name="@mistralai/mistral-vibe", title="Mistral Vibe", version="2.0.2" ) assert response.auth_methods == [] @@ -48,7 +48,7 @@ class TestACPInitialize: ), ) assert response.agent_info == Implementation( - name="@mistralai/mistral-vibe", title="Mistral Vibe", version="2.0.1" + name="@mistralai/mistral-vibe", title="Mistral Vibe", version="2.0.2" ) assert response.auth_methods is not None diff --git a/tests/cli/plan_offer/test_decide_plan_offer.py b/tests/cli/plan_offer/test_decide_plan_offer.py index 0b1d26f..e5b3f55 100644 --- a/tests/cli/plan_offer/test_decide_plan_offer.py +++ b/tests/cli/plan_offer/test_decide_plan_offer.py @@ -5,7 +5,11 @@ import logging import pytest from tests.cli.plan_offer.adapters.fake_whoami_gateway import FakeWhoAmIGateway -from vibe.cli.plan_offer.decide_plan_offer import PlanOfferAction, decide_plan_offer +from vibe.cli.plan_offer.decide_plan_offer import ( + PlanOfferAction, + PlanType, + decide_plan_offer, +) from vibe.cli.plan_offer.ports.whoami_gateway import WhoAmIResponse @@ -18,14 +22,15 @@ async def test_proposes_upgrade_without_call_when_api_key_is_empty() -> None: prompt_switching_to_pro_plan=False, ) ) - action = await decide_plan_offer("", gateway) + action, plan_type = await decide_plan_offer("", gateway) assert action is PlanOfferAction.UPGRADE + assert plan_type is PlanType.FREE assert gateway.calls == [] @pytest.mark.parametrize( - ("response", "expected"), + ("response", "expected_action", "expected_plan_type"), [ ( WhoAmIResponse( @@ -34,6 +39,7 @@ async def test_proposes_upgrade_without_call_when_api_key_is_empty() -> None: prompt_switching_to_pro_plan=False, ), PlanOfferAction.NONE, + PlanType.PRO, ), ( WhoAmIResponse( @@ -42,6 +48,7 @@ async def test_proposes_upgrade_without_call_when_api_key_is_empty() -> None: prompt_switching_to_pro_plan=False, ), PlanOfferAction.UPGRADE, + PlanType.FREE, ), ( WhoAmIResponse( @@ -50,18 +57,22 @@ async def test_proposes_upgrade_without_call_when_api_key_is_empty() -> None: prompt_switching_to_pro_plan=True, ), PlanOfferAction.SWITCH_TO_PRO_KEY, + PlanType.PRO, ), ], ids=["with-a-pro-plan", "without-a-pro-plan", "with-a-non-pro-key"], ) @pytest.mark.asyncio async def test_proposes_an_action_based_on_current_plan_status( - response: WhoAmIResponse, expected: PlanOfferAction + response: WhoAmIResponse, + expected_action: PlanOfferAction, + expected_plan_type: PlanType, ) -> None: gateway = FakeWhoAmIGateway(response) - action = await decide_plan_offer("api-key", gateway) + action, plan_type = await decide_plan_offer("api-key", gateway) - assert action is expected + assert action is expected_action + assert plan_type is expected_plan_type assert gateway.calls == ["api-key"] @@ -75,18 +86,20 @@ async def test_proposes_nothing_when_nothing_is_suggested() -> None: ) ) - action = await decide_plan_offer("api-key", gateway) + action, plan_type = await decide_plan_offer("api-key", gateway) assert action is PlanOfferAction.NONE + assert plan_type is PlanType.UNKNOWN assert gateway.calls == ["api-key"] @pytest.mark.asyncio async def test_proposes_upgrade_when_api_key_is_unauthorized() -> None: gateway = FakeWhoAmIGateway(unauthorized=True) - action = await decide_plan_offer("bad-key", gateway) + action, plan_type = await decide_plan_offer("bad-key", gateway) assert action is PlanOfferAction.UPGRADE + assert plan_type is PlanType.FREE assert gateway.calls == ["bad-key"] @@ -96,8 +109,9 @@ async def test_proposes_none_and_logs_warning_when_gateway_error_occurs( ) -> None: gateway = FakeWhoAmIGateway(error=True) with caplog.at_level(logging.WARNING): - action = await decide_plan_offer("api-key", gateway) + action, plan_type = await decide_plan_offer("api-key", gateway) assert action is PlanOfferAction.NONE + assert plan_type is PlanType.UNKNOWN assert gateway.calls == ["api-key"] assert "Failed to fetch plan status." in caplog.text diff --git a/tests/core/test_config_load_dotenv.py b/tests/core/test_config_load_dotenv.py new file mode 100644 index 0000000..ddd14f3 --- /dev/null +++ b/tests/core/test_config_load_dotenv.py @@ -0,0 +1,61 @@ +from __future__ import annotations + +from pathlib import Path + +from vibe.core.config import load_dotenv_values + + +def _write_env_file(path: Path, content: str) -> None: + path.write_text(content, encoding="utf-8") + + +def test_skips_missing_file(tmp_path: Path) -> None: + environ = {"EXISTING": "1"} + missing_path = tmp_path / "missing.env" + + load_dotenv_values(env_path=missing_path, environ=environ) + + assert environ == {"EXISTING": "1"} + + +def test_sets_and_overrides_values(tmp_path: Path) -> None: + env_path = tmp_path / ".env" + _write_env_file( + env_path, + "\n".join([ + "MISTRAL_API_KEY=new-key", + "HTTPS_PROXY=https://local-proxy:8080", + "OTHER=from-env", + "NEW_KEY=added", + "FOO=replace", + ]) + + "\n", + ) + environ = { + "MISTRAL_API_KEY": "old-key", + "HTTPS_PROXY": "old-https", + "OTHER": "keep", + "FOO": "keep", + } + + load_dotenv_values(env_path=env_path, environ=environ) + + assert environ["MISTRAL_API_KEY"] == "new-key" + assert environ["HTTPS_PROXY"] == "https://local-proxy:8080" + assert environ["OTHER"] == "from-env" + assert environ["NEW_KEY"] == "added" + assert environ["FOO"] == "replace" + + +def test_ignores_empty_values(tmp_path: Path) -> None: + env_path = tmp_path / ".env" + _write_env_file( + env_path, "\n".join(["EMPTY=", "MISTRAL_API_KEY=", "NO_VALUE"]) + "\n" + ) + environ: dict[str, str] = {} + + load_dotenv_values(env_path=env_path, environ=environ) + + assert "EMPTY" not in environ + assert "MISTRAL_API_KEY" not in environ + assert "NO_VALUE" not in environ diff --git a/tests/core/test_config_migration.py b/tests/core/test_config_migration.py deleted file mode 100644 index d2241da..0000000 --- a/tests/core/test_config_migration.py +++ /dev/null @@ -1,51 +0,0 @@ -from __future__ import annotations - -from contextlib import contextmanager -from pathlib import Path -import tomllib - -import tomli_w - -from vibe.core import config -from vibe.core.config import VibeConfig - - -def _restore_dump_config(config_file: Path): - original_dump_config = VibeConfig.dump_config - - def real_dump_config(cls, config_dict: dict) -> None: - try: - with config_file.open("wb") as f: - tomli_w.dump(config_dict, f) - except OSError: - config_file.write_text( - "\n".join( - f"{k} = {v!r}" for k, v in config_dict.items() if v is not None - ), - encoding="utf-8", - ) - - VibeConfig.dump_config = classmethod(real_dump_config) # type: ignore[assignment] - return original_dump_config - - -@contextmanager -def _migrate_config_file(tmp_path: Path, content: str): - config_file = tmp_path / "config.toml" - config_file.write_text(content, encoding="utf-8") - - original_config_file = config.CONFIG_FILE - original_dump_config = _restore_dump_config(config_file) - - try: - config.CONFIG_FILE = config_file - VibeConfig._migrate() - yield config_file - finally: - config.CONFIG_FILE = original_config_file - VibeConfig.dump_config = original_dump_config - - -def _load_migrated_config(config_file: Path) -> dict: - with config_file.open("rb") as f: - return tomllib.load(f) diff --git a/tests/session/test_session_loader.py b/tests/session/test_session_loader.py index 5bc6e1a..cd886a1 100644 --- a/tests/session/test_session_loader.py +++ b/tests/session/test_session_loader.py @@ -206,6 +206,83 @@ class TestSessionLoaderFindLatestSession: assert result is not None assert result == valid_session + def test_find_latest_session_skips_empty_messages_file( + self, session_config: SessionLoggingConfig, create_test_session + ) -> None: + session_dir = Path(session_config.save_dir) + + valid_session = create_test_session(session_dir, "valid123-session") + time.sleep(0.01) + + empty_session = create_test_session(session_dir, "emptymss-session") + (empty_session / "messages.jsonl").write_text("") + + result = SessionLoader.find_latest_session(session_config) + assert result is not None + assert result == valid_session + + def test_find_latest_session_skips_messages_json_not_dict( + self, session_config: SessionLoggingConfig, create_test_session + ) -> None: + session_dir = Path(session_config.save_dir) + + valid_session = create_test_session(session_dir, "valid123-session") + time.sleep(0.01) + + invalid_session = create_test_session(session_dir, "msglist-session") + (invalid_session / "messages.jsonl").write_text("[]\n") + + result = SessionLoader.find_latest_session(session_config) + assert result is not None + assert result == valid_session + + def test_find_latest_session_skips_metadata_json_not_dict( + self, session_config: SessionLoggingConfig, create_test_session + ) -> None: + session_dir = Path(session_config.save_dir) + + valid_session = create_test_session(session_dir, "valid123-session") + time.sleep(0.01) + + invalid_session = create_test_session(session_dir, "metalist-session") + (invalid_session / "meta.json").write_text("[]") + + result = SessionLoader.find_latest_session(session_config) + assert result is not None + assert result == valid_session + + def test_find_latest_session_skips_unreadable_messages_file( + self, session_config: SessionLoggingConfig, create_test_session + ) -> None: + session_dir = Path(session_config.save_dir) + + valid_session = create_test_session(session_dir, "valid123-session") + time.sleep(0.01) + + unreadable_session = create_test_session(session_dir, "unreadab-session") + unreadable_messages = unreadable_session / "messages.jsonl" + unreadable_messages.chmod(0) + + result = SessionLoader.find_latest_session(session_config) + assert result is not None + assert result == valid_session + + def test_find_latest_session_skips_unreadable_metadata_file( + self, session_config: SessionLoggingConfig, create_test_session + ) -> None: + session_dir = Path(session_config.save_dir) + + valid_session = create_test_session(session_dir, "valid123-session") + time.sleep(0.01) + + unreadable_session = create_test_session(session_dir, "unreadab-session") + unreadable_metadata = unreadable_session / "meta.json" + unreadable_metadata.chmod(0) + + result = SessionLoader.find_latest_session(session_config) + assert result is not None + assert result == valid_session + class TestSessionLoaderFindSessionById: def test_find_session_by_id_exact_match( @@ -285,6 +362,27 @@ class TestSessionLoaderFindSessionById: assert result is None +class TestSessionLoaderDoesSessionExist: + def test_does_session_exist_no_messages( + self, session_config: SessionLoggingConfig, create_test_session + ) -> None: + session_dir = Path(session_config.save_dir) + session_folder = create_test_session(session_dir, "test-session-123") + (session_folder / "messages.jsonl").unlink() + + result = SessionLoader.does_session_exist("test-session-123", session_config) + assert result is None + + def test_does_session_exist_success( + self, session_config: SessionLoggingConfig, create_test_session + ) -> None: + session_dir = Path(session_config.save_dir) + session_folder = create_test_session(session_dir, "test-session-123") + + result = SessionLoader.does_session_exist("test-session-123", session_config) + assert result == session_folder + + class TestSessionLoaderLoadSession: def test_load_session_success( self, session_config: SessionLoggingConfig, create_test_session diff --git a/tests/session/test_session_logger.py b/tests/session/test_session_logger.py index daeed43..376fa46 100644 --- a/tests/session/test_session_logger.py +++ b/tests/session/test_session_logger.py @@ -208,8 +208,7 @@ class TestSessionLoggerSaveInteraction: steps=1, session_prompt_tokens=10, session_completion_tokens=20 ) - # Test that save_interaction returns a path when enabled - result = await logger.save_interaction( + await logger.save_interaction( messages=messages, stats=stats, base_config=mock_vibe_config, @@ -217,11 +216,7 @@ class TestSessionLoggerSaveInteraction: agent_profile=mock_agent_profile, ) - # Verify the result - assert result is not None - assert str(logger.session_dir) in result - - # Verify that files were created + # Verify behavior via file system assert logger.session_dir is not None messages_file = logger.session_dir / "messages.jsonl" metadata_file = logger.session_dir / "meta.json" @@ -229,7 +224,6 @@ class TestSessionLoggerSaveInteraction: assert messages_file.exists() assert metadata_file.exists() - # Verify that metadata contains expected data with open(metadata_file) as f: metadata = json.load(f) assert metadata["session_id"] == session_id @@ -261,7 +255,7 @@ class TestSessionLoggerSaveInteraction: steps=1, session_prompt_tokens=10, session_completion_tokens=20 ) - result = await logger.save_interaction( + await logger.save_interaction( messages=messages, stats=stats, base_config=mock_vibe_config, @@ -269,10 +263,9 @@ class TestSessionLoggerSaveInteraction: agent_profile=mock_agent_profile, ) - assert result is not None assert logger.session_dir is not None - metadata_file = logger.session_dir / "meta.json" + assert metadata_file.exists() with open(metadata_file) as f: metadata = json.load(f) assert "system_prompt" in metadata @@ -280,6 +273,7 @@ class TestSessionLoggerSaveInteraction: assert metadata["system_prompt"]["role"] == "system" messages_file = logger.session_dir / "messages.jsonl" + assert messages_file.exists() with open(messages_file) as f: lines = f.readlines() messages_data = [json.loads(line) for line in lines] @@ -330,7 +324,7 @@ class TestSessionLoggerSaveInteraction: steps=2, session_prompt_tokens=20, session_completion_tokens=40 ) - result = await logger.save_interaction( + await logger.save_interaction( messages=all_messages, stats=updated_stats, base_config=mock_vibe_config, @@ -338,17 +332,78 @@ class TestSessionLoggerSaveInteraction: agent_profile=mock_agent_profile, ) - # Verify that the result is not None - assert result is not None - - # Verify that metadata was updated + # Verify behavior via file system: metadata was updated assert logger.session_dir is not None metadata_file = logger.session_dir / "meta.json" + assert metadata_file.exists() with open(metadata_file) as f: metadata = json.load(f) assert metadata["total_messages"] == 4 assert metadata["stats"]["steps"] == updated_stats.steps + messages_file = logger.session_dir / "messages.jsonl" + assert messages_file.exists() + with open(messages_file) as f: + lines = f.readlines() + assert len(lines) == 4 + + @pytest.mark.asyncio + async def test_save_interaction_no_new_messages_is_noop( + self, + session_config: SessionLoggingConfig, + mock_vibe_config: VibeConfig, + mock_tool_manager: ToolManager, + mock_agent_profile: AgentProfile, + ) -> None: + """Test that save_interaction does nothing when there are no new messages.""" + session_id = "test-session-123" + logger = SessionLogger(session_config, session_id) + + messages = [ + LLMMessage(role=Role.system, content="System prompt"), + LLMMessage(role=Role.user, content="Hello"), + LLMMessage(role=Role.assistant, content="Hi there!"), + ] + stats = AgentStats( + steps=1, session_prompt_tokens=10, session_completion_tokens=20 + ) + + await logger.save_interaction( + messages=messages, + stats=stats, + base_config=mock_vibe_config, + tool_manager=mock_tool_manager, + agent_profile=mock_agent_profile, + ) + + assert logger.session_dir is not None + metadata_file = logger.session_dir / "meta.json" + messages_file = logger.session_dir / "messages.jsonl" + + with open(metadata_file) as f: + meta_before = json.load(f) + with open(messages_file) as f: + lines_before = f.readlines() + + # Call again with same messages: no new messages, should be no-op + await logger.save_interaction( + messages=messages, + stats=stats, + base_config=mock_vibe_config, + tool_manager=mock_tool_manager, + agent_profile=mock_agent_profile, + ) + + with open(metadata_file) as f: + meta_after = json.load(f) + with open(messages_file) as f: + lines_after = f.readlines() + + assert len(lines_after) == len(lines_before) == 2 + assert lines_after == lines_before + assert meta_after["total_messages"] == meta_before["total_messages"] == 2 + assert meta_after == meta_before + @pytest.mark.asyncio async def test_save_interaction_no_user_messages( self, @@ -371,7 +426,7 @@ class TestSessionLoggerSaveInteraction: steps=1, session_prompt_tokens=10, session_completion_tokens=20 ) - result = await logger.save_interaction( + await logger.save_interaction( messages=messages, stats=stats, base_config=mock_vibe_config, @@ -379,13 +434,10 @@ class TestSessionLoggerSaveInteraction: agent_profile=mock_agent_profile, ) - # Verify the result - assert result is not None - assert str(logger.session_dir) in result - - # Verify that metadata contains expected data + # Verify behavior via file system assert logger.session_dir is not None metadata_file = logger.session_dir / "meta.json" + assert metadata_file.exists() with open(metadata_file) as f: metadata = json.load(f) assert metadata["session_id"] == session_id @@ -393,6 +445,11 @@ class TestSessionLoggerSaveInteraction: assert metadata["stats"]["steps"] == stats.steps assert metadata["title"] == "Untitled session" + messages_file = logger.session_dir / "messages.jsonl" + assert messages_file.exists() + with open(messages_file) as f: + assert len(f.readlines()) == 1 + @pytest.mark.asyncio async def test_save_interaction_long_user_message( self, @@ -417,7 +474,7 @@ class TestSessionLoggerSaveInteraction: steps=1, session_prompt_tokens=10, session_completion_tokens=20 ) - result = await logger.save_interaction( + await logger.save_interaction( messages=messages, stats=stats, base_config=mock_vibe_config, @@ -425,13 +482,10 @@ class TestSessionLoggerSaveInteraction: agent_profile=mock_agent_profile, ) - # Verify the result - assert result is not None - assert str(logger.session_dir) in result - - # Verify that metadata contains expected data + # Verify behavior via file system assert logger.session_dir is not None metadata_file = logger.session_dir / "meta.json" + assert metadata_file.exists() with open(metadata_file) as f: metadata = json.load(f) assert metadata["session_id"] == session_id @@ -440,6 +494,11 @@ class TestSessionLoggerSaveInteraction: expected_title = long_message[:50] + "…" assert metadata["title"] == expected_title + messages_file = logger.session_dir / "messages.jsonl" + assert messages_file.exists() + with open(messages_file) as f: + assert len(f.readlines()) == 2 + class TestSessionLoggerResetSession: def test_reset_session(self, session_config: SessionLoggingConfig) -> None: diff --git a/tests/snapshots/__snapshots__/test_ui_snapshot_plan_offer/test_snapshot_shows_switch_key_plan_offer.svg b/tests/snapshots/__snapshots__/test_ui_snapshot_plan_offer/test_snapshot_shows_switch_key_plan_offer.svg index 45ff424..d99022b 100644 --- a/tests/snapshots/__snapshots__/test_ui_snapshot_plan_offer/test_snapshot_shows_switch_key_plan_offer.svg +++ b/tests/snapshots/__snapshots__/test_ui_snapshot_plan_offer/test_snapshot_shows_switch_key_plan_offer.svg @@ -35,8 +35,8 @@ .terminal-r1 { fill: #c5c8c6 } .terminal-r2 { fill: #3c3836 } .terminal-r3 { fill: #fbf1c7 } -.terminal-r4 { fill: #fd8019 } -.terminal-r5 { fill: #e5e5e4;font-weight: bold;text-decoration: underline; } +.terminal-r4 { fill: #b3ac90 } +.terminal-r5 { fill: #7b7766 } .terminal-r6 { fill: #a9a9a9 } .terminal-r7 { fill: #a6a087 } .terminal-r8 { fill: #85a598;font-weight: bold } @@ -163,7 +163,7 @@ - + ╭──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╮ @@ -174,10 +174,10 @@ ││ ╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ - -│ -│Switch to your Pro API key -│ +Switch to your Pro API key + + + diff --git a/tests/snapshots/__snapshots__/test_ui_snapshot_plan_offer/test_snapshot_shows_upgrade_plan_offer.svg b/tests/snapshots/__snapshots__/test_ui_snapshot_plan_offer/test_snapshot_shows_upgrade_plan_offer.svg index 493bcdb..0e15b93 100644 --- a/tests/snapshots/__snapshots__/test_ui_snapshot_plan_offer/test_snapshot_shows_upgrade_plan_offer.svg +++ b/tests/snapshots/__snapshots__/test_ui_snapshot_plan_offer/test_snapshot_shows_upgrade_plan_offer.svg @@ -35,8 +35,8 @@ .terminal-r1 { fill: #c5c8c6 } .terminal-r2 { fill: #3c3836 } .terminal-r3 { fill: #fbf1c7 } -.terminal-r4 { fill: #fd8019 } -.terminal-r5 { fill: #e5e5e4;font-weight: bold;text-decoration: underline; } +.terminal-r4 { fill: #b3ac90 } +.terminal-r5 { fill: #7b7766 } .terminal-r6 { fill: #a9a9a9 } .terminal-r7 { fill: #a6a087 } .terminal-r8 { fill: #85a598;font-weight: bold } @@ -163,7 +163,7 @@ - + ╭──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╮ @@ -174,10 +174,10 @@ ││ ╰──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ - -│ -│Upgrade to Pro -│ +Upgrade to Pro + + + diff --git a/tests/snapshots/__snapshots__/test_ui_snapshot_plan_offer/test_snapshot_shows_whats_new_and_plan_offer.svg b/tests/snapshots/__snapshots__/test_ui_snapshot_plan_offer/test_snapshot_shows_whats_new_and_plan_offer.svg index 12b7f75..0fa2c1a 100644 --- a/tests/snapshots/__snapshots__/test_ui_snapshot_plan_offer/test_snapshot_shows_whats_new_and_plan_offer.svg +++ b/tests/snapshots/__snapshots__/test_ui_snapshot_plan_offer/test_snapshot_shows_whats_new_and_plan_offer.svg @@ -38,11 +38,12 @@ .terminal-r4 { fill: #fd8019 } .terminal-r5 { fill: #85a598;font-weight: bold } .terminal-r6 { fill: #aec3bb } -.terminal-r7 { fill: #e5e5e4;font-weight: bold;text-decoration: underline; } -.terminal-r8 { fill: #a9a9a9 } -.terminal-r9 { fill: #a6a087 } -.terminal-r10 { fill: #7e7e7e } -.terminal-r11 { fill: #85a598 } +.terminal-r7 { fill: #b3ac90 } +.terminal-r8 { fill: #7b7766 } +.terminal-r9 { fill: #a9a9a9 } +.terminal-r10 { fill: #a6a087 } +.terminal-r11 { fill: #7e7e7e } +.terminal-r12 { fill: #85a598 } @@ -164,7 +165,7 @@ - + ╭──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╮ @@ -183,9 +184,9 @@ │• Feature 2 │ -│ -│Upgrade to Pro -│ +Upgrade to Pro + + @@ -195,12 +196,12 @@ -⏵ default agent (shift+tab to cycle) +⏵ default agent (shift+tab to cycle) -──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────── ->Ask anything... -──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────── -/test/workdir0% of 200k tokens +──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────── +>Ask anything... +──────────────────────────────────────────────────────────────────────────────────────────────────────────────────────── +/test/workdir0% of 200k tokens diff --git a/tests/test_agent_observer_streaming.py b/tests/test_agent_observer_streaming.py index 6655b19..262f04c 100644 --- a/tests/test_agent_observer_streaming.py +++ b/tests/test_agent_observer_streaming.py @@ -1,9 +1,11 @@ from __future__ import annotations from collections.abc import Callable +from http import HTTPStatus from typing import cast from unittest.mock import AsyncMock +import httpx import pytest from tests.mock.utils import mock_llm_chunk @@ -11,6 +13,7 @@ from tests.stubs.fake_backend import FakeBackend from vibe.core.agent_loop import AgentLoop from vibe.core.agents.models import BuiltinAgentName from vibe.core.config import SessionLoggingConfig, VibeConfig +from vibe.core.llm.exceptions import BackendErrorBuilder from vibe.core.middleware import ( ConversationContext, MiddlewareAction, @@ -25,6 +28,7 @@ from vibe.core.types import ( AssistantEvent, FunctionCall, LLMMessage, + RateLimitError, ReasoningEvent, Role, ToolCall, @@ -369,7 +373,7 @@ async def test_act_handles_user_cancellation_during_streaming() -> None: assert events[-1].skipped is True assert events[-1].skip_reason is not None assert "" in events[-1].skip_reason - assert agent.session_logger.save_interaction.await_count == 1 + assert agent.session_logger.save_interaction.await_count >= 1 @pytest.mark.asyncio @@ -388,6 +392,34 @@ async def test_act_flushes_and_logs_when_streaming_errors(observer_capture) -> N assert agent.session_logger.save_interaction.await_count == 1 +@pytest.mark.asyncio +async def test_rate_limit(observer_capture) -> None: + observed, observer = observer_capture + response = httpx.Response(HTTPStatus.TOO_MANY_REQUESTS) + backend_error = BackendErrorBuilder.build_http_error( + provider="mistral", + endpoint="test", + response=response, + headers=None, + model="test-model", + messages=[], + temperature=0.0, + has_tools=False, + tool_choice=None, + ) + backend = FakeBackend(exception_to_raise=backend_error) + agent = AgentLoop( + make_config(), backend=backend, message_observer=observer, enable_streaming=True + ) + agent.session_logger.save_interaction = AsyncMock(return_value=None) + + with pytest.raises(RateLimitError): + [_ async for _ in agent.act("Trigger rate limit failure while streaming")] + + assert [role for role, _ in observed] == [Role.system, Role.user] + assert agent.session_logger.save_interaction.await_count == 1 + + def _snapshot_events(events: list) -> list[tuple[str, str]]: return [ (type(e).__name__, e.content) diff --git a/tests/test_agents.py b/tests/test_agents.py index 38467c2..bca302f 100644 --- a/tests/test_agents.py +++ b/tests/test_agents.py @@ -1,5 +1,7 @@ from __future__ import annotations +from pathlib import Path + import pytest from tests.mock.utils import mock_llm_chunk @@ -16,6 +18,8 @@ from vibe.core.agents.models import ( _deep_merge, ) from vibe.core.config import SessionLoggingConfig, VibeConfig +from vibe.core.paths.config_paths import ConfigPath +from vibe.core.paths.global_paths import GlobalPath from vibe.core.tools.base import ToolPermission from vibe.core.types import ( FunctionCall, @@ -164,6 +168,43 @@ class TestAgentProfile: } +class TestAgentApplyToConfig: + def test_custom_prompt_found_in_global_when_missing_from_project( + self, tmp_path: Path, monkeypatch: pytest.MonkeyPatch + ) -> None: + """Regression test for https://github.com/mistralai/mistral-vibe/issues/288 + + When a custom prompt .md file is absent from the project-local prompts + directory, the system_prompt property should fall back to the global + ~/.vibe/prompts/ directory and load the file from there. + """ + project_prompts = tmp_path / "project" / ".vibe" / "prompts" + project_prompts.mkdir(parents=True) + + global_prompts = tmp_path / "home" / ".vibe" / "prompts" + global_prompts.mkdir(parents=True) + (global_prompts / "cc.md").write_text("Global custom prompt") + + monkeypatch.setattr( + "vibe.core.config.PROMPTS_DIR", ConfigPath(lambda: project_prompts) + ) + monkeypatch.setattr( + "vibe.core.config.GLOBAL_PROMPTS_DIR", GlobalPath(lambda: global_prompts) + ) + + base = VibeConfig(include_project_context=False, include_prompt_detail=False) + agent = AgentProfile( + name="cc", + display_name="Cc", + description="", + safety=AgentSafety.NEUTRAL, + overrides={"system_prompt_id": "cc"}, + ) + result = agent.apply_to_config(base) + assert result.system_prompt_id == "cc" + assert result.system_prompt == "Global custom prompt" + + class TestAgentProfileOverrides: def test_default_agent_has_no_overrides(self) -> None: assert BUILTIN_AGENTS[BuiltinAgentName.DEFAULT].overrides == {} diff --git a/uv.lock b/uv.lock index d03baf9..033a4f6 100644 --- a/uv.lock +++ b/uv.lock @@ -673,7 +673,7 @@ wheels = [ [[package]] name = "mistral-vibe" -version = "2.0.1" +version = "2.0.2" source = { editable = "." } dependencies = [ { name = "agent-client-protocol" }, diff --git a/vibe/__init__.py b/vibe/__init__.py index 3d78f83..55a122e 100644 --- a/vibe/__init__.py +++ b/vibe/__init__.py @@ -3,4 +3,4 @@ from __future__ import annotations from pathlib import Path VIBE_ROOT = Path(__file__).parent -__version__ = "2.0.1" +__version__ = "2.0.2" diff --git a/vibe/acp/acp_agent_loop.py b/vibe/acp/acp_agent_loop.py index 6c46e5a..8c47cdd 100644 --- a/vibe/acp/acp_agent_loop.py +++ b/vibe/acp/acp_agent_loop.py @@ -65,7 +65,7 @@ from vibe.acp.utils import ( from vibe.core.agent_loop import AgentLoop from vibe.core.agents.models import BuiltinAgentName from vibe.core.autocompletion.path_prompt_adapter import render_path_prompt -from vibe.core.config import MissingAPIKeyError, VibeConfig, load_api_keys_from_env +from vibe.core.config import MissingAPIKeyError, VibeConfig, load_dotenv_values from vibe.core.tools.base import BaseToolConfig, ToolPermission from vibe.core.types import ( ApprovalResponse, @@ -176,7 +176,7 @@ class VibeAcpAgentLoop(AcpAgent): mcp_servers: list[HttpMcpServer | SseMcpServer | McpServerStdio], **kwargs: Any, ) -> NewSessionResponse: - load_api_keys_from_env() + load_dotenv_values() os.chdir(cwd) try: diff --git a/vibe/acp/tools/builtins/bash.py b/vibe/acp/tools/builtins/bash.py index 2c57973..d8c4475 100644 --- a/vibe/acp/tools/builtins/bash.py +++ b/vibe/acp/tools/builtins/bash.py @@ -3,10 +3,8 @@ from __future__ import annotations import asyncio from collections.abc import AsyncGenerator from pathlib import Path -import shlex from acp.schema import ( - EnvVariable, TerminalToolCallContent, ToolCallProgress, ToolCallStart, @@ -40,14 +38,11 @@ class Bash(CoreBashTool, BaseAcpTool[AcpBashState]): timeout = args.timeout or self.config.default_timeout max_bytes = self.config.max_output_bytes - env, command, cmd_args = self._parse_command(args.command) try: terminal_handle = await client.create_terminal( session_id=session_id, - command=command, - args=cmd_args, - env=env, + command=args.command, cwd=str(Path.cwd()), output_byte_limit=max_bytes, ) @@ -84,25 +79,6 @@ class Bash(CoreBashTool, BaseAcpTool[AcpBashState]): except Exception as e: logger.error(f"Failed to release terminal: {e!r}") - def _parse_command( - self, command_str: str - ) -> tuple[list[EnvVariable], str, list[str]]: - parts = shlex.split(command_str) - env: list[EnvVariable] = [] - command: str = "" - args: list[str] = [] - - for part in parts: - if "=" in part and not command: - key, value = part.split("=", 1) - env.append(EnvVariable(name=key, value=value)) - elif not command: - command = part - else: - args.append(part) - - return env, command, args - @classmethod def get_summary(cls, args: BashArgs) -> str: summary = f"{args.command}" diff --git a/vibe/cli/cli.py b/vibe/cli/cli.py index aa21e35..beff89e 100644 --- a/vibe/cli/cli.py +++ b/vibe/cli/cli.py @@ -12,7 +12,7 @@ from vibe.core.config import ( MissingAPIKeyError, MissingPromptFileError, VibeConfig, - load_api_keys_from_env, + load_dotenv_values, ) from vibe.core.paths.config_paths import CONFIG_FILE, HISTORY_FILE from vibe.core.programmatic import run_programmatic @@ -122,7 +122,7 @@ def _load_messages_from_previous_session( def run_cli(args: argparse.Namespace) -> None: - load_api_keys_from_env() + load_dotenv_values() bootstrap_config_files() if args.setup: diff --git a/vibe/cli/plan_offer/decide_plan_offer.py b/vibe/cli/plan_offer/decide_plan_offer.py index 59d0ec2..0895a27 100644 --- a/vibe/cli/plan_offer/decide_plan_offer.py +++ b/vibe/cli/plan_offer/decide_plan_offer.py @@ -29,28 +29,36 @@ ACTION_TO_URL: dict[PlanOfferAction, str] = { } +class PlanType(StrEnum): + FREE = "free" + PRO = "pro" + UNKNOWN = "unknown" + + async def decide_plan_offer( api_key: str | None, gateway: WhoAmIGateway -) -> PlanOfferAction: +) -> tuple[PlanOfferAction, PlanType]: if not api_key: - return PlanOfferAction.UPGRADE + return PlanOfferAction.UPGRADE, PlanType.FREE try: response = await gateway.whoami(api_key) except WhoAmIGatewayUnauthorized: - return PlanOfferAction.UPGRADE + return PlanOfferAction.UPGRADE, PlanType.FREE except WhoAmIGatewayError: logger.warning("Failed to fetch plan status.", exc_info=True) - return PlanOfferAction.NONE - return _action_from_response(response) + return PlanOfferAction.NONE, PlanType.UNKNOWN + return _action_and_plan_from_response(response) -def _action_from_response(response: WhoAmIResponse) -> PlanOfferAction: +def _action_and_plan_from_response( + response: WhoAmIResponse, +) -> tuple[PlanOfferAction, PlanType]: match response: case WhoAmIResponse(is_pro_plan=True): - return PlanOfferAction.NONE + return PlanOfferAction.NONE, PlanType.PRO case WhoAmIResponse(prompt_switching_to_pro_plan=True): - return PlanOfferAction.SWITCH_TO_PRO_KEY + return PlanOfferAction.SWITCH_TO_PRO_KEY, PlanType.PRO case WhoAmIResponse(advertise_pro_plan=True): - return PlanOfferAction.UPGRADE + return PlanOfferAction.UPGRADE, PlanType.FREE case _: - return PlanOfferAction.NONE + return PlanOfferAction.NONE, PlanType.UNKNOWN diff --git a/vibe/cli/textual_ui/app.py b/vibe/cli/textual_ui/app.py index 3c34504..e73b014 100644 --- a/vibe/cli/textual_ui/app.py +++ b/vibe/cli/textual_ui/app.py @@ -23,6 +23,7 @@ from vibe.cli.plan_offer.adapters.http_whoami_gateway import HttpWhoAmIGateway from vibe.cli.plan_offer.decide_plan_offer import ( ACTION_TO_URL, PlanOfferAction, + PlanType, decide_plan_offer, ) from vibe.cli.plan_offer.ports.whoami_gateway import WhoAmIGateway @@ -74,12 +75,19 @@ from vibe.core.agents import AgentProfile from vibe.core.autocompletion.path_prompt_adapter import render_path_prompt from vibe.core.config import VibeConfig from vibe.core.paths.config_paths import HISTORY_FILE +from vibe.core.session.session_loader import SessionLoader from vibe.core.tools.base import ToolPermission from vibe.core.tools.builtins.ask_user_question import ( AskUserQuestionArgs, AskUserQuestionResult, ) -from vibe.core.types import AgentStats, ApprovalResponse, LLMMessage, Role +from vibe.core.types import ( + AgentStats, + ApprovalResponse, + LLMMessage, + RateLimitError, + Role, +) from vibe.core.utils import ( CancellationReason, get_user_cancellation_message, @@ -594,8 +602,16 @@ class VibeApp(App): # noqa: PLR0904 await self._loading_widget.remove() if self.event_handler: self.event_handler.stop_current_tool_call() + + message = str(e) + if isinstance(e, RateLimitError): + if self.plan_type == PlanType.FREE: + message = "Rate limits exceeded. Please wait a moment before trying again, or upgrade to Pro for higher rate limits and uninterrupted access." + else: + message = "Rate limits exceeded. Please wait a moment before trying again." + await self._mount_and_scroll( - ErrorMessage(str(e), collapsed=self._tools_collapsed) + ErrorMessage(message, collapsed=self._tools_collapsed) ) finally: self._agent_running = False @@ -773,6 +789,12 @@ class VibeApp(App): # noqa: PLR0904 return None if not self.agent_loop.session_logger.session_id: return None + session_config = self.agent_loop.session_logger.session_config + session_path = SessionLoader.does_session_exist( + self.agent_loop.session_logger.session_id, session_config + ) + if session_path is None: + return None return self.agent_loop.session_logger.session_id[:8] async def _exit_app(self) -> None: @@ -1023,7 +1045,8 @@ class VibeApp(App): # noqa: PLR0904 async def _maybe_show_plan_offer(self) -> None: if self._plan_offer_shown: return - action = await self._resolve_plan_offer_action() + action, plan_type = await self._resolve_plan_offer_action() + self.plan_type = plan_type if action is PlanOfferAction.NONE: return url = ACTION_TO_URL[action] @@ -1035,12 +1058,12 @@ class VibeApp(App): # noqa: PLR0904 await self._mount_and_scroll(PlanOfferMessage(text)) self._plan_offer_shown = True - async def _resolve_plan_offer_action(self) -> PlanOfferAction: + async def _resolve_plan_offer_action(self) -> tuple[PlanOfferAction, PlanType]: try: active_model = self.config.get_active_model() provider = self.config.get_provider_for_model(active_model) except ValueError: - return PlanOfferAction.NONE + return PlanOfferAction.NONE, PlanType.UNKNOWN api_key_env = provider.api_key_env_var api_key = getenv(api_key_env) if api_key_env else None @@ -1051,7 +1074,7 @@ class VibeApp(App): # noqa: PLR0904 logger.warning( "Plan-offer check failed (%s).", type(exc).__name__, exc_info=True ) - return PlanOfferAction.NONE + return PlanOfferAction.NONE, PlanType.UNKNOWN async def _finalize_current_streaming_message(self) -> None: if self._current_streaming_reasoning is not None: diff --git a/vibe/cli/textual_ui/app.tcss b/vibe/cli/textual_ui/app.tcss index 312d315..066c542 100644 --- a/vibe/cli/textual_ui/app.tcss +++ b/vibe/cli/textual_ui/app.tcss @@ -949,21 +949,24 @@ WelcomeBanner { .whats-new-message { background: $surface; border-left: solid $warning; + margin-bottom: 1; } .plan-offer-message { - background: $surface; - border-left: solid $warning; + background: transparent; + border: none; height: auto; - text-align: center; - content-align: center middle; - padding: 1 2; - margin-top: 1; + padding: 0; + margin-top: 0; + color: $foreground-muted; + text-style: dim; Markdown > *:last-child { margin-bottom: 0; - link-style: bold underline; + link-style: none; + link-color: $foreground-muted; link-background-hover: transparent; + link-style-hover: underline; link-color-hover: $warning; } } diff --git a/vibe/core/agent_loop.py b/vibe/core/agent_loop.py index fea6e5b..78ecdd4 100644 --- a/vibe/core/agent_loop.py +++ b/vibe/core/agent_loop.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio from collections.abc import AsyncGenerator, Callable from enum import StrEnum, auto +from http import HTTPStatus from threading import Thread import time from typing import cast @@ -14,6 +15,7 @@ from vibe.core.agents.manager import AgentManager from vibe.core.agents.models import AgentProfile, BuiltinAgentName from vibe.core.config import VibeConfig from vibe.core.llm.backend.factory import BACKEND_FACTORY +from vibe.core.llm.exceptions import BackendError from vibe.core.llm.format import APIToolFormatHandler, ResolvedMessage, ResolvedToolCall from vibe.core.llm.types import BackendLike from vibe.core.middleware import ( @@ -54,6 +56,7 @@ from vibe.core.types import ( LLMChunk, LLMMessage, LLMUsage, + RateLimitError, ReasoningEvent, Role, SyncApprovalCallback, @@ -95,6 +98,10 @@ class AgentLoopLLMResponseError(AgentLoopError): """Raised when LLM response is malformed or missing expected data.""" +def _should_raise_rate_limit_error(e: Exception) -> bool: + return isinstance(e, BackendError) and e.status == HTTPStatus.TOO_MANY_REQUESTS + + class AgentLoop: def __init__( self, @@ -193,7 +200,18 @@ class AgentLoop: def add_message(self, message: LLMMessage) -> None: self.messages.append(message) - def _flush_new_messages(self) -> None: + async def _save_messages(self) -> None: + await self.session_logger.save_interaction( + self.messages, + self.stats, + self._base_config, + self.tool_manager, + self.agent_profile, + ) + + async def _flush_new_messages(self) -> None: + await self._save_messages() + if not self.message_observer: return @@ -308,12 +326,11 @@ class AgentLoop: if is_user_cancellation_event(event): user_cancelled = True yield event + await self._flush_new_messages() last_message = self.messages[-1] should_break_loop = last_message.role != Role.tool - self._flush_new_messages() - if user_cancelled: return @@ -327,14 +344,7 @@ class AgentLoop: return finally: - self._flush_new_messages() - await self.session_logger.save_interaction( - self.messages, - self.stats, - self._base_config, - self.tool_manager, - self.agent_profile, - ) + await self._flush_new_messages() async def _perform_llm_turn(self) -> AsyncGenerator[BaseEvent, None]: if self.enable_streaming: @@ -593,6 +603,9 @@ class AgentLoop: return LLMChunk(message=processed_message, usage=result.usage) except Exception as e: + if _should_raise_rate_limit_error(e): + raise RateLimitError(provider.name, active_model.name) from e + raise RuntimeError( f"API error from {provider.name} (model: {active_model.name}): {e}" ) from e @@ -642,6 +655,9 @@ class AgentLoop: self.messages.append(chunk_agg.message) except Exception as e: + if _should_raise_rate_limit_error(e): + raise RateLimitError(provider.name, active_model.name) from e + raise RuntimeError( f"API error from {provider.name} (model: {active_model.name}): {e}" ) from e diff --git a/vibe/core/config.py b/vibe/core/config.py index a9f2fbc..1a84d0d 100644 --- a/vibe/core/config.py +++ b/vibe/core/config.py @@ -1,5 +1,6 @@ from __future__ import annotations +from collections.abc import MutableMapping from enum import StrEnum, auto import os from pathlib import Path @@ -19,18 +20,28 @@ from pydantic_settings import ( ) import tomli_w -from vibe.core.paths.config_paths import CONFIG_DIR, CONFIG_FILE, PROMPT_DIR -from vibe.core.paths.global_paths import GLOBAL_ENV_FILE, SESSION_LOG_DIR +from vibe.core.paths.config_paths import CONFIG_DIR, CONFIG_FILE, PROMPTS_DIR +from vibe.core.paths.global_paths import ( + GLOBAL_ENV_FILE, + GLOBAL_PROMPTS_DIR, + SESSION_LOG_DIR, +) from vibe.core.prompts import SystemPrompt from vibe.core.tools.base import BaseToolConfig -def load_api_keys_from_env() -> None: - if GLOBAL_ENV_FILE.path.is_file(): - env_vars = dotenv_values(GLOBAL_ENV_FILE.path) - for key, value in env_vars.items(): - if value: - os.environ.setdefault(key, value) +def load_dotenv_values( + env_path: Path = GLOBAL_ENV_FILE.path, + environ: MutableMapping[str, str] = os.environ, +) -> None: + if not env_path.is_file(): + return + + env_vars = dotenv_values(env_path) + for key, value in env_vars.items(): + if not value: + continue + environ.update({key: value}) class MissingAPIKeyError(RuntimeError): @@ -43,11 +54,17 @@ class MissingAPIKeyError(RuntimeError): class MissingPromptFileError(RuntimeError): - def __init__(self, system_prompt_id: str, prompt_dir: str) -> None: + def __init__( + self, system_prompt_id: str, prompt_dir: str, global_prompt_dir: str + ) -> None: + extra_global_prompt_dir = ( + f" or {global_prompt_dir}" if global_prompt_dir != prompt_dir else "" + ) + super().__init__( f"Invalid system_prompt_id value: '{system_prompt_id}'. " f"Must be one of the available prompts ({', '.join(f'{p.name.lower()}' for p in SystemPrompt)}), " - f"or correspond to a .md file in {prompt_dir}" + f"or correspond to a .md file in {prompt_dir}{extra_global_prompt_dir}" ) self.system_prompt_id = system_prompt_id self.prompt_dir = prompt_dir @@ -392,10 +409,16 @@ class VibeConfig(BaseSettings): except KeyError: pass - custom_sp_path = (PROMPT_DIR.path / self.system_prompt_id).with_suffix(".md") - if not custom_sp_path.is_file(): - raise MissingPromptFileError(self.system_prompt_id, str(PROMPT_DIR.path)) - return custom_sp_path.read_text() + for current_prompt_dir in [PROMPTS_DIR.path, GLOBAL_PROMPTS_DIR.path]: + custom_sp_path = (current_prompt_dir / self.system_prompt_id).with_suffix( + ".md" + ) + if custom_sp_path.is_file(): + return custom_sp_path.read_text() + + raise MissingPromptFileError( + self.system_prompt_id, str(PROMPTS_DIR.path), str(GLOBAL_PROMPTS_DIR.path) + ) def get_active_model(self) -> ModelConfig: for model in self.models: diff --git a/vibe/core/paths/config_paths.py b/vibe/core/paths/config_paths.py index a9690ef..368554e 100644 --- a/vibe/core/paths/config_paths.py +++ b/vibe/core/paths/config_paths.py @@ -62,5 +62,5 @@ def unlock_config_paths() -> None: CONFIG_FILE = ConfigPath(lambda: _resolve_config_path("config.toml", "file")) CONFIG_DIR = ConfigPath(lambda: CONFIG_FILE.path.parent) -PROMPT_DIR = ConfigPath(lambda: _resolve_config_path("prompts", "dir")) +PROMPTS_DIR = ConfigPath(lambda: _resolve_config_path("prompts", "dir")) HISTORY_FILE = ConfigPath(lambda: _resolve_config_path("vibehistory", "file")) diff --git a/vibe/core/paths/global_paths.py b/vibe/core/paths/global_paths.py index 7fa54da..e305552 100644 --- a/vibe/core/paths/global_paths.py +++ b/vibe/core/paths/global_paths.py @@ -31,6 +31,7 @@ GLOBAL_ENV_FILE = GlobalPath(lambda: VIBE_HOME.path / ".env") GLOBAL_TOOLS_DIR = GlobalPath(lambda: VIBE_HOME.path / "tools") GLOBAL_SKILLS_DIR = GlobalPath(lambda: VIBE_HOME.path / "skills") GLOBAL_AGENTS_DIR = GlobalPath(lambda: VIBE_HOME.path / "agents") +GLOBAL_PROMPTS_DIR = GlobalPath(lambda: VIBE_HOME.path / "prompts") SESSION_LOG_DIR = GlobalPath(lambda: VIBE_HOME.path / "logs" / "session") TRUSTED_FOLDERS_FILE = GlobalPath(lambda: VIBE_HOME.path / "trusted_folders.toml") LOG_DIR = GlobalPath(lambda: VIBE_HOME.path / "logs") diff --git a/vibe/core/session/session_loader.py b/vibe/core/session/session_loader.py index 9d239aa..6f9c2e8 100644 --- a/vibe/core/session/session_loader.py +++ b/vibe/core/session/session_loader.py @@ -22,40 +22,48 @@ class SessionLoader: return False try: - with open(metadata_path, encoding="utf-8", errors="ignore") as f: + with metadata_path.open("r", encoding="utf-8", errors="ignore") as f: metadata = json.load(f) - if not isinstance(metadata, dict): - return False + if not isinstance(metadata, dict): + return False - with open(messages_path, encoding="utf-8", errors="ignore") as f: - lines = f.readlines() - if not lines: - return False - messages = [json.loads(line) for line in lines] - if not isinstance(messages, list) or not all( - isinstance(msg, dict) for msg in messages - ): - return False - except json.JSONDecodeError: + with messages_path.open("r", encoding="utf-8", errors="ignore") as f: + has_messages = False + for line in f: + has_messages = True + message = json.loads(line) + if not isinstance(message, dict): + return False + if not has_messages: + return False + except (OSError, UnicodeDecodeError, json.JSONDecodeError): return False return True @staticmethod def latest_session(session_dirs: list[Path]) -> Path | None: - latest_dir = None - latest_mtime = 0 + sessions_with_mtime: list[tuple[Path, float]] = [] for session in session_dirs: - if not SessionLoader._is_valid_session(session): + messages_path = session / MESSAGES_FILENAME + if not messages_path.is_file(): + continue + try: + mtime = messages_path.stat().st_mtime + sessions_with_mtime.append((session, mtime)) + except OSError: continue - messages_path = session / MESSAGES_FILENAME - mtime = messages_path.stat().st_mtime - if mtime > latest_mtime: - latest_mtime = mtime - latest_dir = session + if not sessions_with_mtime: + return None - return latest_dir + sessions_with_mtime.sort(key=lambda x: x[1], reverse=True) + + for session, _mtime in sessions_with_mtime: + if SessionLoader._is_valid_session(session): + return session + + return None @staticmethod def find_latest_session(config: SessionLoggingConfig) -> Path | None: @@ -72,15 +80,32 @@ class SessionLoader: def find_session_by_id( session_id: str, config: SessionLoggingConfig ) -> Path | None: - save_dir = Path(config.save_dir) - if not save_dir.exists(): - return None - - short_id = session_id[:8] - matches = list(save_dir.glob(f"{config.session_prefix}_*_{short_id}")) + matches = SessionLoader._find_session_dirs_by_short_id(session_id, config) return SessionLoader.latest_session(matches) + @staticmethod + def does_session_exist( + session_id: str, config: SessionLoggingConfig + ) -> Path | None: + for session_dir in SessionLoader._find_session_dirs_by_short_id( + session_id, config + ): + if (session_dir / MESSAGES_FILENAME).is_file(): + return session_dir + return None + + @staticmethod + def _find_session_dirs_by_short_id( + session_id: str, config: SessionLoggingConfig + ) -> list[Path]: + save_dir = Path(config.save_dir) + if not save_dir.exists(): + return [] + + short_id = session_id[:8] + return list(save_dir.glob(f"{config.session_prefix}_*_{short_id}")) + @staticmethod def load_session(filepath: Path) -> tuple[list[LLMMessage], dict[str, Any]]: # Load session messages from MESSAGES_FILENAME @@ -94,7 +119,7 @@ class SessionLoader: f"Error reading session messages at {filepath}: {e}" ) from e - if not len(content): + if not content: raise ValueError( f"Session messages file is empty (may have been corrupted by interruption): " f"{filepath}" diff --git a/vibe/core/session/session_logger.py b/vibe/core/session/session_logger.py index c18ddc5..3087f57 100644 --- a/vibe/core/session/session_logger.py +++ b/vibe/core/session/session_logger.py @@ -148,7 +148,7 @@ class SessionLogger: @staticmethod async def persist_metadata(metadata: Any, session_dir: Path) -> None: temp_metadata_filepath = None - metadata_filepath = session_dir / "meta.json" + metadata_filepath = session_dir / METADATA_FILENAME try: async with NamedTemporaryFile( mode="w", @@ -201,12 +201,12 @@ class SessionLogger: base_config: VibeConfig, tool_manager: ToolManager, agent_profile: AgentProfile, - ) -> str | None: + ) -> None: if not self.enabled or self.session_dir is None: - return None + return if self.session_metadata is None: - return None + return # If the session directory does not exist, create it try: @@ -232,14 +232,14 @@ class SessionLogger: ) from e try: + non_system_messages = [m for m in messages if m.role != Role.system] # Append new messages - new_messages = messages[old_total_messages:] + new_messages = non_system_messages[old_total_messages:] - messages_data = [ - m.model_dump(exclude_none=True) - for m in new_messages - if m.role != Role.system - ] + if len(new_messages) == 0: + return + + messages_data = [m.model_dump(exclude_none=True) for m in new_messages] await SessionLogger.persist_messages(messages_data, self.session_dir) # If message update succeeded, write metadata @@ -256,12 +256,12 @@ class SessionLogger: ] title = self._get_title(messages) - if messages[0].role == Role.system: - system_prompt = messages[0].model_dump() - total_messages = len(messages[1:]) - else: - system_prompt = None - total_messages = len(messages) + system_prompt = ( + messages[0].model_dump() + if len(messages) > 0 and messages[0].role == Role.system + else None + ) + total_messages = len(non_system_messages) metadata_dump = { **self.session_metadata.model_dump(), @@ -286,8 +286,6 @@ class SessionLogger: finally: self.cleanup_tmp_files() - return str(self.session_dir) - def reset_session(self, session_id: str) -> None: """Clear existing session info and setup a new session""" if not self.enabled: diff --git a/vibe/core/types.py b/vibe/core/types.py index cda04ab..ce26d05 100644 --- a/vibe/core/types.py +++ b/vibe/core/types.py @@ -382,3 +382,12 @@ type SyncApprovalCallback = Callable[ type ApprovalCallback = AsyncApprovalCallback | SyncApprovalCallback type UserInputCallback = Callable[[BaseModel], Awaitable[BaseModel]] + + +class RateLimitError(Exception): + def __init__(self, provider: str, model: str) -> None: + self.provider = provider + self.model = model + super().__init__( + "Rate limits exceeded. Please wait a moment before trying again." + ) diff --git a/vibe/whats_new.md b/vibe/whats_new.md index e69de29..d38c32b 100644 --- a/vibe/whats_new.md +++ b/vibe/whats_new.md @@ -0,0 +1,5 @@ +# What's New + +- **Config:** Variables defined in the `.env` file in your global `.vibe` folder now override environment variables +- **Sessions:** Fixed message duplication in persisted sessions +- **Resume:** Only shown when a session is available to resume