From f71bfd3b8cdddbb4369c843da9c48d5f6c637ef7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Cl=C3=A9ment=20Drouin?= Date: Wed, 20 May 2026 11:39:59 +0200 Subject: [PATCH] v2.10.1 (#702) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Guillaume LE GOFF Co-authored-by: Michel Thomazo <51709227+michelTho@users.noreply.github.com> Co-authored-by: Pierre Rossinès Co-authored-by: Val <102326092+vdeva@users.noreply.github.com> Co-authored-by: Vincent G <10739306+VinceOPS@users.noreply.github.com> Co-authored-by: Mistral Vibe --- CHANGELOG.md | 14 + distribution/zed/extension.toml | 12 +- pyproject.toml | 2 +- tests/acp/conftest.py | 39 ++ tests/acp/test_acp.py | 8 + tests/acp/test_acp_entrypoint_smoke.py | 115 ++++- tests/acp/test_acp_title.py | 112 +++++ tests/acp/test_authenticate.py | 432 ++++++++++++++++++ tests/acp/test_bash.py | 100 ++-- tests/acp/test_initialize.py | 110 ++++- tests/acp/test_read_file.py | 16 +- tests/acp/test_search_replace.py | 13 +- tests/acp/test_session_auto_title.py | 92 ++++ tests/acp/test_session_set_title.py | 42 +- tests/acp/test_write_file.py | 14 +- tests/autocompletion/test_path_prompt.py | 61 ++- tests/browser_sign_in/stubs.py | 2 +- tests/browser_sign_in/test_browser_sign_in.py | 46 +- .../test_browser_sign_in_http.py | 113 +++-- tests/core/session/test_title_format.py | 74 +++ tests/core/test_config_builder.py | 322 +++++++++++++ tests/core/test_config_orchestrator.py | 86 ++++ tests/core/test_config_toml_end_to_end.py | 71 +++ tests/core/test_user_config_layer.py | 96 ++++ tests/onboarding/test_ui_onboarding.py | 29 +- tests/session/test_session_logger.py | 119 +++++ tests/test_agent_auto_title.py | 83 ++++ uv.lock | 4 +- vibe/__init__.py | 2 +- vibe/acp/acp_agent_loop.py | 320 +++++++++++-- vibe/acp/title.py | 69 +++ vibe/acp/tools/base.py | 42 +- vibe/acp/tools/builtins/bash.py | 17 +- vibe/acp/tools/builtins/read_file.py | 4 +- vibe/acp/tools/builtins/search_replace.py | 6 +- vibe/acp/tools/builtins/write_file.py | 4 +- vibe/acp/tools/events.py | 8 + vibe/cli/textual_ui/app.py | 32 +- vibe/cli/textual_ui/handlers/event_handler.py | 3 + vibe/cli/textual_ui/widgets/mcp_app.py | 16 +- vibe/core/agent_loop.py | 20 +- vibe/core/autocompletion/path_prompt.py | 34 ++ vibe/core/config/builder.py | 122 +++++ vibe/core/config/layers/__init__.py | 5 + vibe/core/config/layers/user.py | 32 ++ vibe/core/config/orchestrator.py | 46 ++ vibe/core/config/schema.py | 20 +- vibe/core/config/types.py | 8 + vibe/core/session/session_logger.py | 35 +- vibe/core/session/title_format.py | 51 +++ vibe/core/types.py | 4 + vibe/setup/auth/__init__.py | 7 +- vibe/setup/auth/api_key_persistence.py | 64 +++ vibe/setup/auth/browser_sign_in.py | 48 +- .../auth/http_browser_sign_in_gateway.py | 35 ++ vibe/setup/onboarding/screens/api_key.py | 63 +-- .../onboarding/screens/browser_sign_in.py | 8 +- 57 files changed, 2950 insertions(+), 402 deletions(-) create mode 100644 tests/acp/test_acp_title.py create mode 100644 tests/acp/test_authenticate.py create mode 100644 tests/acp/test_session_auto_title.py create mode 100644 tests/core/session/test_title_format.py create mode 100644 tests/core/test_config_builder.py create mode 100644 tests/core/test_config_orchestrator.py create mode 100644 tests/core/test_config_toml_end_to_end.py create mode 100644 tests/core/test_user_config_layer.py create mode 100644 tests/test_agent_auto_title.py create mode 100644 vibe/acp/title.py create mode 100644 vibe/acp/tools/events.py create mode 100644 vibe/core/config/builder.py create mode 100644 vibe/core/config/layers/__init__.py create mode 100644 vibe/core/config/layers/user.py create mode 100644 vibe/core/config/orchestrator.py create mode 100644 vibe/core/config/types.py create mode 100644 vibe/core/session/title_format.py create mode 100644 vibe/setup/auth/api_key_persistence.py diff --git a/CHANGELOG.md b/CHANGELOG.md index e0bc77f..7df74d9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,20 @@ 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.10.1] - 2026-05-20 + +### Added + +- Pretty session titles that format `@mention` syntax for human-readable display +- Auto-emit `SessionInfoUpdate` on the first prompt so IDE session pickers reflect the title immediately +- End-to-end layered config read path with `TomlFileLayer`, `ConfigBuilder`, and `ConfigOrchestrator` + +### Fixed + +- MCP menu no longer shows `D`/`E` shortcuts in the detail view when there are no tools to enable or disable +- Parallel `bash` tool calls in ACP each render their own live terminal instead of racing on shared state + + ## [2.10.0] - 2026-05-19 ### Added diff --git a/distribution/zed/extension.toml b/distribution/zed/extension.toml index 61cb0ff..336d461 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.10.0" +version = "2.10.1" schema_version = 1 authors = ["Mistral AI"] repository = "https://github.com/mistralai/mistral-vibe" @@ -11,21 +11,21 @@ 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.10.0/vibe-acp-darwin-aarch64-2.10.0.zip" +archive = "https://github.com/mistralai/mistral-vibe/releases/download/v2.10.1/vibe-acp-darwin-aarch64-2.10.1.zip" cmd = "./vibe-acp" [agent_servers.mistral-vibe.targets.darwin-x86_64] -archive = "https://github.com/mistralai/mistral-vibe/releases/download/v2.10.0/vibe-acp-darwin-x86_64-2.10.0.zip" +archive = "https://github.com/mistralai/mistral-vibe/releases/download/v2.10.1/vibe-acp-darwin-x86_64-2.10.1.zip" cmd = "./vibe-acp" [agent_servers.mistral-vibe.targets.linux-aarch64] -archive = "https://github.com/mistralai/mistral-vibe/releases/download/v2.10.0/vibe-acp-linux-aarch64-2.10.0.zip" +archive = "https://github.com/mistralai/mistral-vibe/releases/download/v2.10.1/vibe-acp-linux-aarch64-2.10.1.zip" cmd = "./vibe-acp" [agent_servers.mistral-vibe.targets.linux-x86_64] -archive = "https://github.com/mistralai/mistral-vibe/releases/download/v2.10.0/vibe-acp-linux-x86_64-2.10.0.zip" +archive = "https://github.com/mistralai/mistral-vibe/releases/download/v2.10.1/vibe-acp-linux-x86_64-2.10.1.zip" cmd = "./vibe-acp" [agent_servers.mistral-vibe.targets.windows-x86_64] -archive = "https://github.com/mistralai/mistral-vibe/releases/download/v2.10.0/vibe-acp-windows-x86_64-2.10.0.zip" +archive = "https://github.com/mistralai/mistral-vibe/releases/download/v2.10.1/vibe-acp-windows-x86_64-2.10.1.zip" cmd = "./vibe-acp.exe" diff --git a/pyproject.toml b/pyproject.toml index 492aa8c..a218d1d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "mistral-vibe" -version = "2.10.0" +version = "2.10.1" description = "Minimal CLI coding agent by Mistral" readme = "README.md" requires-python = ">=3.12" diff --git a/tests/acp/conftest.py b/tests/acp/conftest.py index e03ad3f..45e7662 100644 --- a/tests/acp/conftest.py +++ b/tests/acp/conftest.py @@ -7,10 +7,12 @@ from unittest.mock import patch import pytest +from tests.conftest import build_test_vibe_config from tests.stubs.fake_backend import FakeBackend from tests.stubs.fake_client import FakeClient from vibe.acp.acp_agent_loop import VibeAcpAgentLoop from vibe.core.agent_loop import AgentLoop +from vibe.core.config import ModelConfig, SessionLoggingConfig from vibe.core.types import LLMChunk, LLMMessage, LLMUsage, Role @@ -45,6 +47,43 @@ def acp_agent_loop(backend: FakeBackend) -> VibeAcpAgentLoop: return _create_acp_agent() +@pytest.fixture +def acp_agent_with_session_config( + backend: FakeBackend, temp_session_dir: Path, monkeypatch: pytest.MonkeyPatch +) -> tuple[VibeAcpAgentLoop, FakeClient]: + session_config = SessionLoggingConfig( + save_dir=str(temp_session_dir), session_prefix="session", enabled=True + ) + config = build_test_vibe_config( + active_model="devstral-latest", + models=[ + ModelConfig( + name="devstral-latest", provider="mistral", alias="devstral-latest" + ) + ], + session_logging=session_config, + ) + + class PatchedAgentLoop(AgentLoop): + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **{**kwargs, "backend": backend}) + self._base_config = config + self.agent_manager.invalidate_config() + + monkeypatch.setattr("vibe.acp.acp_agent_loop.AgentLoop", PatchedAgentLoop) + monkeypatch.setattr(VibeAcpAgentLoop, "_load_config", lambda self: config) + monkeypatch.setattr( + "vibe.acp.acp_agent_loop.VibeConfig.load", lambda *args, **kwargs: config + ) + + vibe_acp_agent = VibeAcpAgentLoop() + client = FakeClient() + vibe_acp_agent.on_connect(client) + client.on_connect(vibe_acp_agent) + + return vibe_acp_agent, client + + @pytest.fixture def temp_session_dir(tmp_path: Path) -> Path: session_dir = tmp_path / "sessions" diff --git a/tests/acp/test_acp.py b/tests/acp/test_acp.py index 32fe5ac..aad994e 100644 --- a/tests/acp/test_acp.py +++ b/tests/acp/test_acp.py @@ -463,6 +463,14 @@ class TestSessionUpdates: == "available_commands_update" ) + title_response_text = await read_response(process) + assert title_response_text is not None + title_response = UpdateJsonRpcNotification.model_validate( + json.loads(title_response_text) + ) + assert title_response.params is not None + assert title_response.params.update.session_update == "session_info_update" + text_response = await read_response(process) assert text_response is not None response = UpdateJsonRpcNotification.model_validate( diff --git a/tests/acp/test_acp_entrypoint_smoke.py b/tests/acp/test_acp_entrypoint_smoke.py index e6a11cc..1230ea6 100644 --- a/tests/acp/test_acp_entrypoint_smoke.py +++ b/tests/acp/test_acp_entrypoint_smoke.py @@ -16,6 +16,11 @@ import pytest from tests import TESTS_ROOT from tests.e2e.common import ansi_tolerant_pattern +BROWSER_AUTH_NAME = "Sign in through Mistral AI Studio" +BROWSER_AUTH_DESCRIPTION = ( + "Sign into Mistral Vibe through your Mistral AI Studio account." +) + class _AcpSmokeClient(Client): def on_connect(self, conn: Any) -> None: @@ -96,15 +101,22 @@ async def _terminate_process(proc: asyncio.subprocess.Process) -> None: await proc.wait() -def _build_env(vibe_home_dir: Path, *, include_api_key: bool) -> dict[str, str]: +def _build_env( + vibe_home_dir: Path, *, include_api_key: bool, enable_browser_sign_in: bool = False +) -> dict[str, str]: env = os.environ.copy() env["PYTHONUNBUFFERED"] = "1" env["VIBE_HOME"] = str(vibe_home_dir) vibe_home_dir.mkdir(parents=True, exist_ok=True) + config_lines = ["enable_telemetry = false"] + if enable_browser_sign_in: + config_lines.append("enable_experimental_browser_sign_in = true") config_file = vibe_home_dir / "config.toml" if not config_file.exists(): - config_file.write_text("enable_telemetry = false\n") + config_file.write_text("\n".join(config_lines) + "\n") + elif enable_browser_sign_in: + config_file.write_text(config_file.read_text() + config_lines[-1] + "\n") if include_api_key: env["MISTRAL_API_KEY"] = "mock" @@ -114,17 +126,33 @@ def _build_env(vibe_home_dir: Path, *, include_api_key: bool) -> dict[str, str]: return env -def _build_client_capabilities(*, terminal_auth: bool = False) -> ClientCapabilities: - if not terminal_auth: +def _build_client_capabilities( + *, terminal_auth: bool = False, delegated_browser_auth: bool = False +) -> ClientCapabilities: + if not terminal_auth and not delegated_browser_auth: return ClientCapabilities() - return ClientCapabilities(field_meta={"terminal-auth": True}) + field_meta: dict[str, bool] = {} + if terminal_auth: + field_meta["terminal-auth"] = True + if delegated_browser_auth: + field_meta["browser-auth-delegated"] = True + return ClientCapabilities(field_meta=field_meta) async def _connect_and_initialize( - *, vibe_home_dir: Path, include_api_key: bool, terminal_auth: bool = False + *, + vibe_home_dir: Path, + include_api_key: bool, + enable_browser_sign_in: bool = False, + terminal_auth: bool = False, + delegated_browser_auth: bool = False, ) -> tuple[asyncio.subprocess.Process, Any, Any]: - env = _build_env(vibe_home_dir, include_api_key=include_api_key) + env = _build_env( + vibe_home_dir, + include_api_key=include_api_key, + enable_browser_sign_in=enable_browser_sign_in, + ) proc = await _spawn_vibe_acp(env) try: @@ -136,7 +164,8 @@ async def _connect_and_initialize( conn.initialize( protocol_version=PROTOCOL_VERSION, client_capabilities=_build_client_capabilities( - terminal_auth=terminal_auth + terminal_auth=terminal_auth, + delegated_browser_auth=delegated_browser_auth, ), client_info=Implementation( name="pytest-smoke", title="Pytest Smoke", version="0.0.0" @@ -188,18 +217,71 @@ async def test_vibe_acp_bootstraps_default_files(vibe_home_dir: Path) -> None: @pytest.mark.asyncio -async def test_vibe_acp_initialize_exposes_terminal_auth_when_supported( - vibe_home_dir: Path, -) -> None: - proc, initialize_response, conn = await _connect_and_initialize( - vibe_home_dir=vibe_home_dir, include_api_key=True, terminal_auth=True +async def test_vibe_acp_initialize_exposes_browser_auth(vibe_home_dir: Path) -> None: + proc, initialize_response, _conn = await _connect_and_initialize( + vibe_home_dir=vibe_home_dir, include_api_key=True, enable_browser_sign_in=True ) try: assert initialize_response.auth_methods is not None assert len(initialize_response.auth_methods) == 1 - auth_method = initialize_response.auth_methods[0] + assert auth_method.id == "browser-auth" + assert auth_method.name == BROWSER_AUTH_NAME + assert auth_method.description == BROWSER_AUTH_DESCRIPTION + finally: + await _terminate_process(proc) + + +@pytest.mark.asyncio +async def test_vibe_acp_initialize_exposes_delegated_browser_auth_when_supported( + vibe_home_dir: Path, +) -> None: + proc, initialize_response, _conn = await _connect_and_initialize( + vibe_home_dir=vibe_home_dir, + include_api_key=True, + enable_browser_sign_in=True, + delegated_browser_auth=True, + ) + + try: + assert initialize_response.auth_methods is not None + assert len(initialize_response.auth_methods) == 2 + + browser_auth_method = initialize_response.auth_methods[0] + assert browser_auth_method.id == "browser-auth" + assert browser_auth_method.name == BROWSER_AUTH_NAME + assert browser_auth_method.description == BROWSER_AUTH_DESCRIPTION + + delegated_browser_auth_method = initialize_response.auth_methods[1] + assert delegated_browser_auth_method.id == "browser-auth-delegated" + assert delegated_browser_auth_method.name == BROWSER_AUTH_NAME + assert delegated_browser_auth_method.description == BROWSER_AUTH_DESCRIPTION + finally: + await _terminate_process(proc) + + +@pytest.mark.asyncio +async def test_vibe_acp_initialize_exposes_terminal_auth_when_supported( + vibe_home_dir: Path, +) -> None: + proc, initialize_response, _conn = await _connect_and_initialize( + vibe_home_dir=vibe_home_dir, + include_api_key=True, + enable_browser_sign_in=True, + terminal_auth=True, + ) + + try: + assert initialize_response.auth_methods is not None + assert len(initialize_response.auth_methods) == 2 + + browser_auth_method = initialize_response.auth_methods[0] + assert browser_auth_method.id == "browser-auth" + assert browser_auth_method.name == BROWSER_AUTH_NAME + assert browser_auth_method.description == BROWSER_AUTH_DESCRIPTION + + auth_method = initialize_response.auth_methods[1] assert auth_method.id == "vibe-setup" assert auth_method.field_meta is not None @@ -207,11 +289,12 @@ async def test_vibe_acp_initialize_exposes_terminal_auth_when_supported( assert terminal_auth["label"] == "Mistral Vibe Setup" assert terminal_auth["command"] assert terminal_auth["args"] + assert terminal_auth["args"][-1:] == ["--setup"] finally: await _terminate_process(proc) -@pytest.mark.timeout(15) +@pytest.mark.timeout(30) def test_vibe_acp_setup_shows_onboarding_and_exits_on_cancel( vibe_home_dir: Path, ) -> None: @@ -231,7 +314,7 @@ def test_vibe_acp_setup_shows_onboarding_and_exits_on_cancel( child.logfile_read = captured try: - child.expect(ansi_tolerant_pattern("Welcome to Mistral Vibe"), timeout=10) + child.expect(ansi_tolerant_pattern("Welcome to"), timeout=15) child.sendcontrol("c") child.expect(pexpect.EOF, timeout=10) finally: diff --git a/tests/acp/test_acp_title.py b/tests/acp/test_acp_title.py new file mode 100644 index 0000000..b4d08dd --- /dev/null +++ b/tests/acp/test_acp_title.py @@ -0,0 +1,112 @@ +from __future__ import annotations + +from acp.schema import ( + EmbeddedResourceContentBlock, + ResourceContentBlock, + TextContentBlock, + TextResourceContents, +) + +from vibe.acp.title import acp_blocks_to_title_segments +from vibe.core.session.title_format import MentionSegment, TextSegment + + +class TestAcpBlocksToTitleSegments: + def test_text_block(self) -> None: + segments = acp_blocks_to_title_segments([ + TextContentBlock(type="text", text="hello") + ]) + assert segments == [TextSegment(text="hello")] + + def test_resource_link_with_name(self) -> None: + block = ResourceContentBlock( + type="resource_link", uri="file:///abs/path/foo.py", name="foo.py" + ) + segments = acp_blocks_to_title_segments([block]) + assert segments == [MentionSegment(name="foo.py")] + + def test_resource_link_with_line_range_fragment(self) -> None: + block = ResourceContentBlock( + type="resource_link", uri="file:///abs/path/foo.py#L9-L27", name="foo.py" + ) + segments = acp_blocks_to_title_segments([block]) + assert segments == [MentionSegment(name="foo.py", start_line=9, end_line=27)] + + def test_resource_link_with_single_line_fragment(self) -> None: + block = ResourceContentBlock( + type="resource_link", uri="file:///abs/path/foo.py#L42", name="foo.py" + ) + segments = acp_blocks_to_title_segments([block]) + assert segments == [MentionSegment(name="foo.py", start_line=42, end_line=None)] + + def test_resource_link_falls_back_to_uri_basename_when_no_name(self) -> None: + block = ResourceContentBlock( + type="resource_link", uri="file:///abs/path/bar.txt", name="" + ) + segments = acp_blocks_to_title_segments([block]) + assert segments == [MentionSegment(name="bar.txt")] + + def test_embedded_resource_block_derives_basename_from_uri(self) -> None: + block = EmbeddedResourceContentBlock( + type="resource", + resource=TextResourceContents( + uri="file:///abs/path/script.sh", text="echo hi" + ), + ) + segments = acp_blocks_to_title_segments([block]) + assert segments == [MentionSegment(name="script.sh")] + + def test_embedded_resource_with_line_range_fragment(self) -> None: + block = EmbeddedResourceContentBlock( + type="resource", + resource=TextResourceContents( + uri="file:///abs/path/main.py#L1-L20", text="content" + ), + ) + segments = acp_blocks_to_title_segments([block]) + assert segments == [MentionSegment(name="main.py", start_line=1, end_line=20)] + + def test_automatic_field_meta_skips_block(self) -> None: + block = EmbeddedResourceContentBlock( + type="resource", + field_meta={"automatic": True}, + resource=TextResourceContents(uri="file:///abs/path/auto.py", text="..."), + ) + segments = acp_blocks_to_title_segments([block]) + assert segments == [] + + def test_automatic_field_meta_skips_resource_link(self) -> None: + block = ResourceContentBlock( + type="resource_link", + uri="file:///abs/path/auto.py", + name="auto.py", + field_meta={"automatic": True}, + ) + segments = acp_blocks_to_title_segments([block]) + assert segments == [] + + def test_preserves_order(self) -> None: + blocks = [ + TextContentBlock(type="text", text="Look at "), + ResourceContentBlock( + type="resource_link", uri="file:///a/foo.py", name="foo.py" + ), + TextContentBlock(type="text", text=" and "), + ResourceContentBlock( + type="resource_link", uri="file:///a/bar.py", name="bar.py" + ), + ] + segments = acp_blocks_to_title_segments(blocks) + assert segments == [ + TextSegment(text="Look at "), + MentionSegment(name="foo.py"), + TextSegment(text=" and "), + MentionSegment(name="bar.py"), + ] + + def test_unknown_fragment_format_ignored(self) -> None: + block = ResourceContentBlock( + type="resource_link", uri="file:///abs/path/foo.py#section1", name="foo.py" + ) + segments = acp_blocks_to_title_segments([block]) + assert segments == [MentionSegment(name="foo.py")] diff --git a/tests/acp/test_authenticate.py b/tests/acp/test_authenticate.py new file mode 100644 index 0000000..dc7b9fa --- /dev/null +++ b/tests/acp/test_authenticate.py @@ -0,0 +1,432 @@ +from __future__ import annotations + +from datetime import UTC, datetime +from typing import Any + +import pytest + +from vibe.acp.acp_agent_loop import VibeAcpAgentLoop +from vibe.acp.exceptions import InternalError, InvalidRequestError +from vibe.core.config import ProviderConfig +from vibe.core.types import Backend +from vibe.setup.auth import ( + BrowserSignInAttempt, + BrowserSignInError, + BrowserSignInErrorCode, +) +from vibe.setup.onboarding.context import OnboardingContext + + +def build_browser_sign_in_attempt( + process_id: str = "process-123", +) -> BrowserSignInAttempt: + return BrowserSignInAttempt( + process_id=process_id, + sign_in_url=f"https://console.mistral.ai/vibe/sign-in/{process_id}", + poll_url=f"https://console.mistral.ai/api/vibe/sign-in/{process_id}", + expires_at=datetime(2026, 4, 23, 12, 0, tzinfo=UTC), + code_verifier="secret-code-verifier", + ) + + +def build_mistral_provider( + *, + api_key_env_var: str = "MISTRAL_API_KEY", + browser_auth_base_url: str = "https://console.mistral.ai", + browser_auth_api_base_url: str = "https://console.mistral.ai/api", +) -> ProviderConfig: + return ProviderConfig( + name="mistral", + api_base="https://api.mistral.ai/v1", + api_key_env_var=api_key_env_var, + browser_auth_base_url=browser_auth_base_url, + browser_auth_api_base_url=browser_auth_api_base_url, + backend=Backend.MISTRAL, + ) + + +def build_unsupported_provider() -> ProviderConfig: + return ProviderConfig( + name="llamacpp", + api_base="http://127.0.0.1:8080/v1", + api_key_env_var="LLAMACPP_API_KEY", + backend=Backend.GENERIC, + ) + + +class MutableOnboardingContextLoader: + def __init__( + self, provider: ProviderConfig, *, enable_browser_sign_in: bool = True + ) -> None: + self.provider = provider + self.enable_browser_sign_in = enable_browser_sign_in + + def __call__(self) -> OnboardingContext: + return OnboardingContext( + provider=self.provider, + enable_experimental_browser_sign_in=self.enable_browser_sign_in, + ) + + +class FakeBrowserSignInService: + def __init__( + self, + *, + attempt: BrowserSignInAttempt | None = None, + api_key: str = "api-key", + authenticate_error: BrowserSignInError | None = None, + start_error: BrowserSignInError | None = None, + complete_errors: list[BrowserSignInError] | None = None, + complete_error: BrowserSignInError | None = None, + ) -> None: + self.attempt = attempt or build_browser_sign_in_attempt() + self.api_key = api_key + self.authenticate_error = authenticate_error + self.start_error = start_error + self.complete_errors = list(complete_errors or []) + if complete_error is not None: + self.complete_errors.append(complete_error) + self.close_count = 0 + + async def authenticate(self) -> str: + if self.authenticate_error is not None: + raise self.authenticate_error + return self.api_key + + async def start_attempt(self) -> BrowserSignInAttempt: + if self.start_error is not None: + raise self.start_error + return self.attempt + + async def complete_attempt(self, attempt: BrowserSignInAttempt) -> str: + if self.complete_errors: + raise self.complete_errors.pop(0) + if attempt != self.attempt: + raise AssertionError("Unexpected browser sign-in attempt.") + return self.api_key + + async def aclose(self) -> None: + self.close_count += 1 + + +class InMemoryApiKeyPersister: + def __init__(self, result: str = "completed") -> None: + self.result = result + self.saved: list[tuple[ProviderConfig, str]] = [] + + def persist(self, provider: ProviderConfig, api_key: str) -> str: + self.saved.append((provider, api_key)) + return self.result + + +def build_acp_agent( + *, + provider: ProviderConfig | None = None, + enable_browser_sign_in: bool = True, + browser_sign_in: FakeBrowserSignInService | None = None, + api_key_persister: InMemoryApiKeyPersister | None = None, +) -> tuple[VibeAcpAgentLoop, MutableOnboardingContextLoader, InMemoryApiKeyPersister]: + provider = provider or build_mistral_provider() + browser_sign_in = browser_sign_in or FakeBrowserSignInService() + api_key_persister = api_key_persister or InMemoryApiKeyPersister() + context_loader = MutableOnboardingContextLoader( + provider, enable_browser_sign_in=enable_browser_sign_in + ) + + return ( + VibeAcpAgentLoop( + onboarding_context_loader=context_loader, + browser_sign_in_service_factory=lambda _provider: browser_sign_in, + api_key_persister=api_key_persister.persist, + ), + context_loader, + api_key_persister, + ) + + +def require_auth_meta(response: Any, method_id: str) -> dict[str, Any]: + assert response is not None + assert response.field_meta is not None + meta = response.field_meta[method_id] + assert isinstance(meta, dict) + return meta + + +class TestACPAuthenticate: + @pytest.mark.asyncio + async def test_authenticate_completes_browser_sign_in_and_persists_api_key( + self, + ) -> None: + provider = build_mistral_provider() + browser_sign_in = FakeBrowserSignInService(api_key="api-key") + acp_agent_loop, _, api_key_persister = build_acp_agent( + provider=provider, browser_sign_in=browser_sign_in + ) + + response = await acp_agent_loop.authenticate("browser-auth") + + assert require_auth_meta(response, "browser-auth") == { + "persistResult": "completed", + "status": "completed", + } + assert api_key_persister.saved == [(provider, "api-key")] + assert browser_sign_in.close_count == 1 + + @pytest.mark.asyncio + async def test_authenticate_starts_delegated_browser_sign_in(self) -> None: + attempt = build_browser_sign_in_attempt() + browser_sign_in = FakeBrowserSignInService(attempt=attempt) + acp_agent_loop, _, api_key_persister = build_acp_agent( + browser_sign_in=browser_sign_in + ) + + response = await acp_agent_loop.authenticate("browser-auth-delegated") + + assert require_auth_meta(response, "browser-auth-delegated") == { + "attemptId": "process-123", + "expiresAt": "2026-04-23T12:00:00Z", + "signInUrl": "https://console.mistral.ai/vibe/sign-in/process-123", + } + assert api_key_persister.saved == [] + assert browser_sign_in.close_count == 1 + + @pytest.mark.asyncio + async def test_authenticate_rejects_unsupported_method( + self, acp_agent_loop: VibeAcpAgentLoop + ) -> None: + with pytest.raises( + InvalidRequestError, match="Unsupported auth method: vibe-setup" + ): + await acp_agent_loop.authenticate("vibe-setup") + + @pytest.mark.asyncio + async def test_authenticate_rejects_browser_sign_in_when_unavailable(self) -> None: + acp_agent_loop = VibeAcpAgentLoop( + onboarding_context_loader=lambda: OnboardingContext( + provider=build_unsupported_provider() + ) + ) + + with pytest.raises( + InvalidRequestError, + match="Browser sign-in is not available for the configured provider.", + ): + await acp_agent_loop.authenticate("browser-auth") + + @pytest.mark.asyncio + async def test_authenticate_rejects_browser_sign_in_when_flag_disabled( + self, + ) -> None: + browser_sign_in = FakeBrowserSignInService(api_key="api-key") + acp_agent_loop, _, api_key_persister = build_acp_agent( + browser_sign_in=browser_sign_in, enable_browser_sign_in=False + ) + + with pytest.raises( + InvalidRequestError, + match="Browser sign-in is not available for the configured provider.", + ): + await acp_agent_loop.authenticate("browser-auth") + + assert api_key_persister.saved == [] + assert browser_sign_in.close_count == 0 + + @pytest.mark.asyncio + async def test_authenticate_rejects_delegated_browser_sign_in_when_flag_disabled( + self, + ) -> None: + browser_sign_in = FakeBrowserSignInService() + acp_agent_loop, _, api_key_persister = build_acp_agent( + browser_sign_in=browser_sign_in, enable_browser_sign_in=False + ) + + with pytest.raises( + InvalidRequestError, + match="Browser sign-in is not available for the configured provider.", + ): + await acp_agent_loop.authenticate("browser-auth-delegated") + + assert api_key_persister.saved == [] + assert browser_sign_in.close_count == 0 + + @pytest.mark.asyncio + async def test_authenticate_surfaces_start_failures(self) -> None: + browser_sign_in = FakeBrowserSignInService( + authenticate_error=BrowserSignInError( + "Failed to start browser sign-in.", + code=BrowserSignInErrorCode.START_FAILED, + ) + ) + acp_agent_loop, _, _ = build_acp_agent(browser_sign_in=browser_sign_in) + + with pytest.raises(InternalError, match="Failed to start browser sign-in."): + await acp_agent_loop.authenticate("browser-auth") + + assert browser_sign_in.close_count == 1 + + @pytest.mark.asyncio + async def test_authenticate_completes_delegated_browser_sign_in_and_persists_api_key( + self, + ) -> None: + provider = build_mistral_provider() + attempt = build_browser_sign_in_attempt() + browser_sign_in = FakeBrowserSignInService(attempt=attempt, api_key="api-key") + acp_agent_loop, _, api_key_persister = build_acp_agent( + provider=provider, browser_sign_in=browser_sign_in + ) + start_response = await acp_agent_loop.authenticate("browser-auth-delegated") + attempt_id = require_auth_meta(start_response, "browser-auth-delegated")[ + "attemptId" + ] + + response = await acp_agent_loop.authenticate( + "browser-auth-delegated", action="complete", attemptId=attempt_id + ) + + assert require_auth_meta(response, "browser-auth-delegated") == { + "attemptId": "process-123", + "persistResult": "completed", + "status": "completed", + } + assert api_key_persister.saved == [(provider, "api-key")] + assert browser_sign_in.close_count == 2 + + @pytest.mark.asyncio + async def test_authenticate_delegated_completion_uses_provider_captured_at_start( + self, + ) -> None: + start_provider = build_mistral_provider() + current_provider = build_mistral_provider( + api_key_env_var="OTHER_API_KEY", + browser_auth_base_url="https://example.com", + browser_auth_api_base_url="https://example.com/api", + ) + browser_sign_in = FakeBrowserSignInService(api_key="api-key") + acp_agent_loop, context_loader, api_key_persister = build_acp_agent( + provider=start_provider, browser_sign_in=browser_sign_in + ) + start_response = await acp_agent_loop.authenticate("browser-auth-delegated") + attempt_id = require_auth_meta(start_response, "browser-auth-delegated")[ + "attemptId" + ] + context_loader.provider = current_provider + + await acp_agent_loop.authenticate( + "browser-auth-delegated", action="complete", attemptId=attempt_id + ) + + assert api_key_persister.saved == [(start_provider, "api-key")] + + @pytest.mark.asyncio + async def test_authenticate_delegated_completion_allows_flag_disabled_after_start( + self, + ) -> None: + provider = build_mistral_provider() + browser_sign_in = FakeBrowserSignInService(api_key="api-key") + acp_agent_loop, context_loader, api_key_persister = build_acp_agent( + provider=provider, browser_sign_in=browser_sign_in + ) + start_response = await acp_agent_loop.authenticate("browser-auth-delegated") + attempt_id = require_auth_meta(start_response, "browser-auth-delegated")[ + "attemptId" + ] + context_loader.enable_browser_sign_in = False + + await acp_agent_loop.authenticate( + "browser-auth-delegated", action="complete", attemptId=attempt_id + ) + + assert api_key_persister.saved == [(provider, "api-key")] + + @pytest.mark.asyncio + async def test_authenticate_delegated_completion_requires_attempt_id( + self, acp_agent_loop: VibeAcpAgentLoop + ) -> None: + with pytest.raises( + InvalidRequestError, match="Missing browser sign-in attempt ID." + ): + await acp_agent_loop.authenticate( + "browser-auth-delegated", action="complete" + ) + + @pytest.mark.asyncio + async def test_authenticate_delegated_completion_rejects_unknown_attempt_id( + self, + ) -> None: + acp_agent_loop, _, _ = build_acp_agent() + + with pytest.raises( + InvalidRequestError, match="Unknown browser sign-in attempt: process-123" + ): + await acp_agent_loop.authenticate( + "browser-auth-delegated", action="complete", attemptId="process-123" + ) + + @pytest.mark.asyncio + async def test_authenticate_delegated_completion_surfaces_browser_sign_in_failures( + self, + ) -> None: + browser_sign_in = FakeBrowserSignInService( + complete_error=BrowserSignInError( + "Browser sign-in timed out.", code=BrowserSignInErrorCode.TIMED_OUT + ) + ) + acp_agent_loop, _, api_key_persister = build_acp_agent( + browser_sign_in=browser_sign_in + ) + start_response = await acp_agent_loop.authenticate("browser-auth-delegated") + attempt_id = require_auth_meta(start_response, "browser-auth-delegated")[ + "attemptId" + ] + + with pytest.raises(InvalidRequestError, match="Browser sign-in timed out."): + await acp_agent_loop.authenticate( + "browser-auth-delegated", action="complete", attemptId=attempt_id + ) + + assert api_key_persister.saved == [] + assert browser_sign_in.close_count == 2 + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "error_code", + [BrowserSignInErrorCode.EXCHANGE_FAILED, BrowserSignInErrorCode.POLL_FAILED], + ) + async def test_authenticate_delegated_completion_keeps_retryable_attempts( + self, error_code: BrowserSignInErrorCode + ) -> None: + provider = build_mistral_provider() + browser_sign_in = FakeBrowserSignInService( + api_key="api-key", + complete_errors=[ + BrowserSignInError( + "Transient browser sign-in failure.", code=error_code + ) + ], + ) + acp_agent_loop, _, api_key_persister = build_acp_agent( + provider=provider, browser_sign_in=browser_sign_in + ) + start_response = await acp_agent_loop.authenticate("browser-auth-delegated") + attempt_id = require_auth_meta(start_response, "browser-auth-delegated")[ + "attemptId" + ] + + with pytest.raises( + InvalidRequestError, match="Transient browser sign-in failure." + ): + await acp_agent_loop.authenticate( + "browser-auth-delegated", action="complete", attemptId=attempt_id + ) + + retry_response = await acp_agent_loop.authenticate( + "browser-auth-delegated", action="complete", attemptId=attempt_id + ) + + assert require_auth_meta(retry_response, "browser-auth-delegated") == { + "attemptId": attempt_id, + "persistResult": "completed", + "status": "completed", + } + assert api_key_persister.saved == [(provider, "api-key")] + assert browser_sign_in.close_count == 3 diff --git a/tests/acp/test_bash.py b/tests/acp/test_bash.py index c09b232..ac19b52 100644 --- a/tests/acp/test_bash.py +++ b/tests/acp/test_bash.py @@ -8,7 +8,8 @@ import pytest from tests.mock.utils import collect_result from vibe.acp.tools.builtins.bash import AcpBashState, Bash -from vibe.core.tools.base import ToolError +from vibe.acp.tools.events import ToolTerminalOpenedEvent +from vibe.core.tools.base import InvokeContext, ToolError from vibe.core.tools.builtins.bash import BashArgs, BashResult, BashToolConfig @@ -105,9 +106,7 @@ def acp_bash_tool(mock_client: MockClient) -> Bash: config = BashToolConfig() # Use model_construct to bypass Pydantic validation for testing state = AcpBashState.model_construct( - client=mock_client, - session_id="test_session_123", - tool_call_id="test_tool_call_456", + client=mock_client, session_id="test_session_123" ) return Bash(config_getter=lambda: config, state=state) @@ -156,7 +155,7 @@ class TestAcpBashExecution: tool = Bash( config_getter=lambda: BashToolConfig(), state=AcpBashState.model_construct( - client=mock_client, session_id="test_session", tool_call_id="test_call" + client=mock_client, session_id="test_session" ), ) @@ -176,7 +175,7 @@ class TestAcpBashExecution: tool = Bash( config_getter=lambda: BashToolConfig(), state=AcpBashState.model_construct( - client=mock_client, session_id="test_session", tool_call_id="test_call" + client=mock_client, session_id="test_session" ), ) @@ -196,7 +195,7 @@ class TestAcpBashExecution: tool = Bash( config_getter=lambda: BashToolConfig(), state=AcpBashState.model_construct( - client=mock_client, session_id="test_session", tool_call_id="test_call" + client=mock_client, session_id="test_session" ), ) @@ -213,9 +212,7 @@ class TestAcpBashExecution: async def test_run_without_client(self) -> None: tool = Bash( config_getter=lambda: BashToolConfig(), - state=AcpBashState.model_construct( - client=None, session_id="test_session", tool_call_id="test_call" - ), + state=AcpBashState.model_construct(client=None, session_id="test_session"), ) args = BashArgs(command="test") @@ -232,9 +229,7 @@ class TestAcpBashExecution: mock_client = MockClient() tool = Bash( config_getter=lambda: BashToolConfig(), - state=AcpBashState.model_construct( - client=mock_client, session_id=None, tool_call_id="test_call" - ), + state=AcpBashState.model_construct(client=mock_client, session_id=None), ) args = BashArgs(command="test") @@ -256,7 +251,7 @@ class TestAcpBashExecution: tool = Bash( config_getter=lambda: BashToolConfig(), state=AcpBashState.model_construct( - client=mock_client, session_id="test_session", tool_call_id="test_call" + client=mock_client, session_id="test_session" ), ) @@ -283,7 +278,7 @@ class TestAcpBashTimeout: tool = Bash( config_getter=lambda: BashToolConfig(default_timeout=30), state=AcpBashState.model_construct( - client=mock_client, session_id="test_session", tool_call_id="test_call" + client=mock_client, session_id="test_session" ), ) @@ -312,7 +307,7 @@ class TestAcpBashTimeout: tool = Bash( config_getter=lambda: BashToolConfig(), state=AcpBashState.model_construct( - client=mock_client, session_id="test_session", tool_call_id="test_call" + client=mock_client, session_id="test_session" ), ) @@ -324,61 +319,56 @@ class TestAcpBashTimeout: assert str(exc_info.value) == "Command timed out after 1s: 'slow_command'" -class TestAcpBashEmbedding: +class TestAcpBashTerminalOpenedEvent: @pytest.mark.asyncio - async def test_run_with_embedding(self, mock_client: MockClient) -> None: - tool = Bash( - config_getter=lambda: BashToolConfig(), - state=AcpBashState.model_construct( - client=mock_client, session_id="test_session", tool_call_id="test_call" - ), - ) - - args = BashArgs(command="test") - await collect_result(tool.run(args)) - - assert mock_client._session_update_called - - @pytest.mark.asyncio - async def test_run_embedding_without_tool_call_id( + async def test_run_yields_terminal_opened_event( self, mock_client: MockClient ) -> None: tool = Bash( config_getter=lambda: BashToolConfig(), state=AcpBashState.model_construct( - client=mock_client, session_id="test_session", tool_call_id=None + client=mock_client, session_id="test_session" ), ) args = BashArgs(command="test") - await collect_result(tool.run(args)) + events: list[ToolTerminalOpenedEvent] = [] + async for item in tool.run(args, InvokeContext(tool_call_id="test_call")): + if isinstance(item, ToolTerminalOpenedEvent): + events.append(item) - # Embedding should be skipped when tool_call_id is None - assert not mock_client._session_update_called + assert len(events) == 1 + assert events[0].terminal_id == mock_client._terminal_handle.id + assert events[0].tool_call_id == "test_call" + assert events[0].tool_name == "bash" + +class TestAcpBashConcurrentInvocations: @pytest.mark.asyncio - async def test_run_embedding_handles_exception( - self, mock_client: MockClient - ) -> None: - # Make session_update raise an exception - async def failing_session_update(session_id: str, update, **kwargs) -> None: - raise RuntimeError("Session update failed") - - mock_client.session_update = failing_session_update - + async def test_concurrent_invocations_yield_distinct_tool_call_ids(self) -> None: + mock_client = MockClient(MockTerminalHandle(terminal_id="t", wait_delay=0.05)) tool = Bash( config_getter=lambda: BashToolConfig(), state=AcpBashState.model_construct( - client=mock_client, session_id="test_session", tool_call_id="test_call" + client=mock_client, session_id="test_session" ), ) - args = BashArgs(command="test") - # Should not raise, embedding failure is silently ignored - result = await collect_result(tool.run(args)) + async def run_and_collect_ids(tool_call_id: str) -> list[str]: + ids: list[str] = [] + async for item in tool.run( + BashArgs(command="echo hi"), InvokeContext(tool_call_id=tool_call_id) + ): + if isinstance(item, ToolTerminalOpenedEvent): + ids.append(item.tool_call_id) + return ids - assert result is not None - assert result.stdout == "test output" + results = await asyncio.gather( + run_and_collect_ids("T1"), run_and_collect_ids("T2") + ) + + assert results[0] == ["T1"] + assert results[1] == ["T2"] class TestAcpBashConfig: @@ -395,7 +385,7 @@ class TestAcpBashConfig: tool = Bash( config_getter=lambda: BashToolConfig(default_timeout=30), state=AcpBashState.model_construct( - client=mock_client, session_id="test_session", tool_call_id="test_call" + client=mock_client, session_id="test_session" ), ) @@ -425,7 +415,7 @@ class TestAcpBashCleanup: tool = Bash( config_getter=lambda: BashToolConfig(), state=AcpBashState.model_construct( - client=mock_client, session_id="test_session", tool_call_id="test_call" + client=mock_client, session_id="test_session" ), ) @@ -457,7 +447,7 @@ class TestAcpBashCleanup: tool = Bash( config_getter=lambda: BashToolConfig(), state=AcpBashState.model_construct( - client=mock_client, session_id="test_session", tool_call_id="test_call" + client=mock_client, session_id="test_session" ), ) @@ -483,7 +473,7 @@ class TestAcpBashCleanup: tool = Bash( config_getter=lambda: BashToolConfig(), state=AcpBashState.model_construct( - client=mock_client, session_id="test_session", tool_call_id="test_call" + client=mock_client, session_id="test_session" ), ) diff --git a/tests/acp/test_initialize.py b/tests/acp/test_initialize.py index 9995965..629a127 100644 --- a/tests/acp/test_initialize.py +++ b/tests/acp/test_initialize.py @@ -14,11 +14,42 @@ from acp.schema import ( import pytest from vibe.acp.acp_agent_loop import VibeAcpAgentLoop +from vibe.core.config import ProviderConfig +from vibe.core.types import Backend +from vibe.setup.onboarding.context import OnboardingContext + +BROWSER_AUTH_NAME = "Sign in through Mistral AI Studio" +BROWSER_AUTH_DESCRIPTION = ( + "Sign into Mistral Vibe through your Mistral AI Studio account." +) + + +def build_mistral_provider() -> ProviderConfig: + return ProviderConfig( + name="mistral", + api_base="https://api.mistral.ai/v1", + api_key_env_var="MISTRAL_API_KEY", + browser_auth_base_url="https://console.mistral.ai", + browser_auth_api_base_url="https://console.mistral.ai/api", + backend=Backend.MISTRAL, + ) + + +def build_acp_agent_loop( + *, provider: ProviderConfig | None = None, enable_browser_sign_in: bool = True +) -> VibeAcpAgentLoop: + return VibeAcpAgentLoop( + onboarding_context_loader=lambda: OnboardingContext( + provider=provider or build_mistral_provider(), + enable_experimental_browser_sign_in=enable_browser_sign_in, + ) + ) class TestACPInitialize: @pytest.mark.asyncio - async def test_initialize(self, acp_agent_loop: VibeAcpAgentLoop) -> None: + async def test_initialize(self) -> None: + acp_agent_loop = build_acp_agent_loop() response = await acp_agent_loop.initialize(protocol_version=PROTOCOL_VERSION) assert response.protocol_version == PROTOCOL_VERSION @@ -34,16 +65,20 @@ class TestACPInitialize: ), ) assert response.agent_info == Implementation( - name="@mistralai/mistral-vibe", title="Mistral Vibe", version="2.10.0" + name="@mistralai/mistral-vibe", title="Mistral Vibe", version="2.10.1" ) - assert response.auth_methods == [] + assert response.auth_methods is not None + assert len(response.auth_methods) == 1 + auth_method = response.auth_methods[0] + assert auth_method.id == "browser-auth" + assert auth_method.name == BROWSER_AUTH_NAME + assert auth_method.description == BROWSER_AUTH_DESCRIPTION @pytest.mark.asyncio - async def test_initialize_with_terminal_auth( - self, acp_agent_loop: VibeAcpAgentLoop - ) -> None: + async def test_initialize_with_terminal_auth(self) -> None: """Test initialize with terminal-auth capabilities to check it was included.""" + acp_agent_loop = build_acp_agent_loop() client_capabilities = ClientCapabilities(field_meta={"terminal-auth": True}) response = await acp_agent_loop.initialize( protocol_version=PROTOCOL_VERSION, client_capabilities=client_capabilities @@ -62,12 +97,18 @@ class TestACPInitialize: ), ) assert response.agent_info == Implementation( - name="@mistralai/mistral-vibe", title="Mistral Vibe", version="2.10.0" + name="@mistralai/mistral-vibe", title="Mistral Vibe", version="2.10.1" ) assert response.auth_methods is not None - assert len(response.auth_methods) == 1 - auth_method = response.auth_methods[0] + assert len(response.auth_methods) == 2 + + browser_auth_method = response.auth_methods[0] + assert browser_auth_method.id == "browser-auth" + assert browser_auth_method.name == BROWSER_AUTH_NAME + assert browser_auth_method.description == BROWSER_AUTH_DESCRIPTION + + auth_method = response.auth_methods[1] assert auth_method.id == "vibe-setup" assert auth_method.name == "Register your API Key" assert auth_method.description == "Register your API Key inside Mistral Vibe" @@ -76,4 +117,55 @@ class TestACPInitialize: terminal_auth_meta = auth_method.field_meta["terminal-auth"] assert "command" in terminal_auth_meta assert "args" in terminal_auth_meta + assert terminal_auth_meta["args"][-1:] == ["--setup"] assert terminal_auth_meta["label"] == "Mistral Vibe Setup" + + @pytest.mark.asyncio + async def test_initialize_with_delegated_browser_auth(self) -> None: + acp_agent_loop = build_acp_agent_loop() + client_capabilities = ClientCapabilities( + field_meta={"browser-auth-delegated": True} + ) + response = await acp_agent_loop.initialize( + protocol_version=PROTOCOL_VERSION, client_capabilities=client_capabilities + ) + + assert response.auth_methods is not None + assert len(response.auth_methods) == 2 + + browser_auth_method = response.auth_methods[0] + assert browser_auth_method.id == "browser-auth" + assert browser_auth_method.name == BROWSER_AUTH_NAME + assert browser_auth_method.description == BROWSER_AUTH_DESCRIPTION + + delegated_browser_auth_method = response.auth_methods[1] + assert delegated_browser_auth_method.id == "browser-auth-delegated" + assert delegated_browser_auth_method.name == BROWSER_AUTH_NAME + assert delegated_browser_auth_method.description == BROWSER_AUTH_DESCRIPTION + + @pytest.mark.asyncio + async def test_initialize_omits_browser_auth_when_experimental_flag_disabled( + self, + ) -> None: + acp_agent_loop = build_acp_agent_loop(enable_browser_sign_in=False) + + response = await acp_agent_loop.initialize(protocol_version=PROTOCOL_VERSION) + + assert response.auth_methods == [] + + @pytest.mark.asyncio + async def test_initialize_omits_browser_auth_when_provider_unsupported( + self, + ) -> None: + acp_agent_loop = build_acp_agent_loop( + provider=ProviderConfig( + name="llamacpp", + api_base="http://127.0.0.1:8080/v1", + api_key_env_var="LLAMACPP_API_KEY", + backend=Backend.GENERIC, + ) + ) + + response = await acp_agent_loop.initialize(protocol_version=PROTOCOL_VERSION) + + assert response.auth_methods == [] diff --git a/tests/acp/test_read_file.py b/tests/acp/test_read_file.py index 3c704e4..1d2bab4 100644 --- a/tests/acp/test_read_file.py +++ b/tests/acp/test_read_file.py @@ -74,7 +74,6 @@ def acp_read_file_tool( state = AcpReadFileState.model_construct( client=mock_client, # type: ignore[arg-type] session_id="test_session_123", - tool_call_id="test_tool_call_456", ) return ReadFile(config_getter=lambda: config, state=state) @@ -99,7 +98,6 @@ class TestAcpReadFileExecution: assert result.content == "line 1\nline 2\nline 3" assert result.lines_read == 3 assert mock_client._read_text_file_called - assert mock_client._session_update_called # Verify read_text_file was called correctly params = mock_client._last_read_params @@ -118,7 +116,7 @@ class TestAcpReadFileExecution: tool = ReadFile( config_getter=lambda: ReadFileToolConfig(), state=AcpReadFileState.model_construct( - client=mock_client, session_id="test_session", tool_call_id="test_call" + client=mock_client, session_id="test_session" ), ) @@ -141,7 +139,7 @@ class TestAcpReadFileExecution: tool = ReadFile( config_getter=lambda: ReadFileToolConfig(), state=AcpReadFileState.model_construct( - client=mock_client, session_id="test_session", tool_call_id="test_call" + client=mock_client, session_id="test_session" ), ) @@ -164,7 +162,7 @@ class TestAcpReadFileExecution: tool = ReadFile( config_getter=lambda: ReadFileToolConfig(), state=AcpReadFileState.model_construct( - client=mock_client, session_id="test_session", tool_call_id="test_call" + client=mock_client, session_id="test_session" ), ) @@ -189,7 +187,7 @@ class TestAcpReadFileExecution: tool = ReadFile( config_getter=lambda: ReadFileToolConfig(), state=AcpReadFileState.model_construct( - client=mock_client, session_id="test_session", tool_call_id="test_call" + client=mock_client, session_id="test_session" ), ) @@ -209,7 +207,7 @@ class TestAcpReadFileExecution: tool = ReadFile( config_getter=lambda: ReadFileToolConfig(), state=AcpReadFileState.model_construct( - client=None, session_id="test_session", tool_call_id="test_call" + client=None, session_id="test_session" ), ) @@ -232,9 +230,7 @@ class TestAcpReadFileExecution: mock_client = MockClient() tool = ReadFile( config_getter=lambda: ReadFileToolConfig(), - state=AcpReadFileState.model_construct( - client=mock_client, session_id=None, tool_call_id="test_call" - ), + state=AcpReadFileState.model_construct(client=mock_client, session_id=None), ) args = ReadFileArgs(path=str(test_file)) diff --git a/tests/acp/test_search_replace.py b/tests/acp/test_search_replace.py index a90b644..d3e35fe 100644 --- a/tests/acp/test_search_replace.py +++ b/tests/acp/test_search_replace.py @@ -81,9 +81,7 @@ def acp_search_replace_tool( monkeypatch.chdir(tmp_path) config = SearchReplaceConfig() state = AcpSearchReplaceState.model_construct( - client=mock_client, - session_id="test_session_123", - tool_call_id="test_tool_call_456", + client=mock_client, session_id="test_session_123" ) return SearchReplace(config_getter=lambda: config, state=state) @@ -116,7 +114,6 @@ class TestAcpSearchReplaceExecution: assert result.blocks_applied == 1 assert mock_client._read_text_file_called assert mock_client._write_text_file_called - assert mock_client._session_update_called # Verify read_text_file was called correctly read_params = mock_client._last_read_params @@ -141,7 +138,7 @@ class TestAcpSearchReplaceExecution: tool = SearchReplace( config_getter=lambda: config, state=AcpSearchReplaceState.model_construct( - client=mock_client, session_id="test_session", tool_call_id="test_call" + client=mock_client, session_id="test_session" ), ) @@ -171,7 +168,7 @@ class TestAcpSearchReplaceExecution: tool = SearchReplace( config_getter=lambda: SearchReplaceConfig(), state=AcpSearchReplaceState.model_construct( - client=mock_client, session_id="test_session", tool_call_id="test_call" + client=mock_client, session_id="test_session" ), ) @@ -202,7 +199,7 @@ class TestAcpSearchReplaceExecution: tool = SearchReplace( config_getter=lambda: SearchReplaceConfig(), state=AcpSearchReplaceState.model_construct( - client=mock_client, session_id="test_session", tool_call_id="test_call" + client=mock_client, session_id="test_session" ), ) @@ -245,7 +242,7 @@ class TestAcpSearchReplaceExecution: tool = SearchReplace( config_getter=lambda: SearchReplaceConfig(), state=AcpSearchReplaceState.model_construct( - client=client, session_id=session_id, tool_call_id="test_call" + client=client, session_id=session_id ), ) diff --git a/tests/acp/test_session_auto_title.py b/tests/acp/test_session_auto_title.py new file mode 100644 index 0000000..3fc0726 --- /dev/null +++ b/tests/acp/test_session_auto_title.py @@ -0,0 +1,92 @@ +from __future__ import annotations + +from pathlib import Path + +from acp.schema import ResourceContentBlock, SessionInfoUpdate, TextContentBlock +import pytest + +from tests.stubs.fake_client import FakeClient +from vibe.acp.acp_agent_loop import VibeAcpAgentLoop + + +def _info_updates(client: FakeClient) -> list: + return [ + notification.update + for notification in client._session_updates + if isinstance(notification.update, SessionInfoUpdate) + ] + + +class TestAcpAutoTitleOnPrompt: + @pytest.mark.asyncio + async def test_emits_session_info_update_on_first_prompt( + self, acp_agent_with_session_config: tuple[VibeAcpAgentLoop, FakeClient] + ) -> None: + acp_agent, client = acp_agent_with_session_config + + new_session = await acp_agent.new_session(cwd=str(Path.cwd()), mcp_servers=[]) + assert new_session is not None + + await acp_agent.prompt( + session_id=new_session.session_id, + prompt=[ + TextContentBlock(type="text", text="Refactor "), + ResourceContentBlock( + type="resource_link", uri="file:///abs/auth.py", name="auth.py" + ), + TextContentBlock(type="text", text=" please"), + ], + ) + + updates = _info_updates(client) + assert len(updates) == 1 + assert updates[0].title == "Refactor @auth.py please" + assert updates[0].updated_at is None + + @pytest.mark.asyncio + async def test_no_event_on_second_prompt( + self, acp_agent_with_session_config: tuple[VibeAcpAgentLoop, FakeClient] + ) -> None: + acp_agent, client = acp_agent_with_session_config + + new_session = await acp_agent.new_session(cwd=str(Path.cwd()), mcp_servers=[]) + assert new_session is not None + + await acp_agent.prompt( + session_id=new_session.session_id, + prompt=[TextContentBlock(type="text", text="first")], + ) + assert len(_info_updates(client)) == 1 + + await acp_agent.prompt( + session_id=new_session.session_id, + prompt=[TextContentBlock(type="text", text="second")], + ) + + assert len(_info_updates(client)) == 1 + + @pytest.mark.asyncio + async def test_skips_automatic_resource_in_title( + self, acp_agent_with_session_config: tuple[VibeAcpAgentLoop, FakeClient] + ) -> None: + acp_agent, client = acp_agent_with_session_config + + new_session = await acp_agent.new_session(cwd=str(Path.cwd()), mcp_servers=[]) + assert new_session is not None + + await acp_agent.prompt( + session_id=new_session.session_id, + prompt=[ + TextContentBlock(type="text", text="What does this do?"), + ResourceContentBlock( + type="resource_link", + uri="file:///abs/open_in_editor.py", + name="open_in_editor.py", + field_meta={"automatic": True}, + ), + ], + ) + + updates = _info_updates(client) + assert len(updates) == 1 + assert updates[0].title == "What does this do?" diff --git a/tests/acp/test_session_set_title.py b/tests/acp/test_session_set_title.py index 7109de1..73c7e6d 100644 --- a/tests/acp/test_session_set_title.py +++ b/tests/acp/test_session_set_title.py @@ -7,50 +7,10 @@ from acp.schema import SessionInfoUpdate import pytest import tomli_w -from tests.conftest import build_test_vibe_config, get_base_config -from tests.stubs.fake_backend import FakeBackend +from tests.conftest import get_base_config from tests.stubs.fake_client import FakeClient from vibe.acp.acp_agent_loop import VibeAcpAgentLoop from vibe.acp.exceptions import InternalError, InvalidRequestError, SessionNotFoundError -from vibe.core.agent_loop import AgentLoop -from vibe.core.config import ModelConfig, SessionLoggingConfig - - -@pytest.fixture -def acp_agent_with_session_config( - backend: FakeBackend, temp_session_dir: Path, monkeypatch: pytest.MonkeyPatch -) -> tuple[VibeAcpAgentLoop, FakeClient]: - session_config = SessionLoggingConfig( - save_dir=str(temp_session_dir), session_prefix="session", enabled=True - ) - config = build_test_vibe_config( - active_model="devstral-latest", - models=[ - ModelConfig( - name="devstral-latest", provider="mistral", alias="devstral-latest" - ) - ], - session_logging=session_config, - ) - - class PatchedAgentLoop(AgentLoop): - def __init__(self, *args, **kwargs) -> None: - super().__init__(*args, **{**kwargs, "backend": backend}) - self._base_config = config - self.agent_manager.invalidate_config() - - monkeypatch.setattr("vibe.acp.acp_agent_loop.AgentLoop", PatchedAgentLoop) - monkeypatch.setattr(VibeAcpAgentLoop, "_load_config", lambda self: config) - monkeypatch.setattr( - "vibe.acp.acp_agent_loop.VibeConfig.load", lambda *args, **kwargs: config - ) - - vibe_acp_agent = VibeAcpAgentLoop() - client = FakeClient() - vibe_acp_agent.on_connect(client) - client.on_connect(vibe_acp_agent) - - return vibe_acp_agent, client class TestSessionSetTitle: diff --git a/tests/acp/test_write_file.py b/tests/acp/test_write_file.py index f627f7c..f9e666c 100644 --- a/tests/acp/test_write_file.py +++ b/tests/acp/test_write_file.py @@ -54,9 +54,7 @@ def acp_write_file_tool( monkeypatch.chdir(tmp_path) config = WriteFileConfig() state = AcpWriteFileState.model_construct( - client=mock_client, - session_id="test_session_123", - tool_call_id="test_tool_call_456", + client=mock_client, session_id="test_session_123" ) return WriteFile(config_getter=lambda: config, state=state) @@ -81,7 +79,6 @@ class TestAcpWriteFileExecution: assert result.bytes_written == len(b"Hello, world!") assert result.file_existed is False assert mock_client._write_text_file_called - assert mock_client._session_update_called # Verify write_text_file was called correctly params = mock_client._last_write_params @@ -97,7 +94,7 @@ class TestAcpWriteFileExecution: tool = WriteFile( config_getter=lambda: WriteFileConfig(), state=AcpWriteFileState.model_construct( - client=mock_client, session_id="test_session", tool_call_id="test_call" + client=mock_client, session_id="test_session" ), ) @@ -114,7 +111,6 @@ class TestAcpWriteFileExecution: assert result.bytes_written == len(b"New content") assert result.file_existed is True assert mock_client._write_text_file_called - assert mock_client._session_update_called # Verify write_text_file was called correctly params = mock_client._last_write_params @@ -132,7 +128,7 @@ class TestAcpWriteFileExecution: tool = WriteFile( config_getter=lambda: WriteFileConfig(), state=AcpWriteFileState.model_construct( - client=mock_client, session_id="test_session", tool_call_id="test_call" + client=mock_client, session_id="test_session" ), ) @@ -151,7 +147,7 @@ class TestAcpWriteFileExecution: tool = WriteFile( config_getter=lambda: WriteFileConfig(), state=AcpWriteFileState.model_construct( - client=None, session_id="test_session", tool_call_id="test_call" + client=None, session_id="test_session" ), ) @@ -173,7 +169,7 @@ class TestAcpWriteFileExecution: tool = WriteFile( config_getter=lambda: WriteFileConfig(), state=AcpWriteFileState.model_construct( - client=mock_client, session_id=None, tool_call_id="test_call" + client=mock_client, session_id=None ), ) diff --git a/tests/autocompletion/test_path_prompt.py b/tests/autocompletion/test_path_prompt.py index 01e48fd..18a5967 100644 --- a/tests/autocompletion/test_path_prompt.py +++ b/tests/autocompletion/test_path_prompt.py @@ -2,7 +2,11 @@ from __future__ import annotations from pathlib import Path -from vibe.core.autocompletion.path_prompt import build_path_prompt_payload +from vibe.core.autocompletion.path_prompt import ( + build_path_prompt_payload, + build_title_segments, +) +from vibe.core.session.title_format import MentionSegment, TextSegment def test_deduplicates_same_file_mentioned_twice(tmp_path: Path) -> None: @@ -16,3 +20,58 @@ def test_deduplicates_same_file_mentioned_twice(tmp_path: Path) -> None: assert len(payload.resources) == 1 assert payload.resources[0].path == readme assert len(payload.all_resources) == 2 + + +class TestBuildTitleSegments: + def test_empty_message(self) -> None: + assert build_title_segments("") == [] + + def test_plain_text_no_mentions(self) -> None: + segments = build_title_segments("hello world") + assert segments == [TextSegment(text="hello world")] + + def test_matched_file_mention_uses_basename(self, tmp_path: Path) -> None: + nested = tmp_path / "src" / "auth" + nested.mkdir(parents=True) + target = nested / "foo.py" + target.write_text("x", encoding="utf-8") + + segments = build_title_segments( + "Refactor @src/auth/foo.py please", base_dir=tmp_path + ) + assert segments == [ + TextSegment(text="Refactor "), + MentionSegment(name="foo.py"), + TextSegment(text=" please"), + ] + + def test_unmatched_mention_stays_as_text(self, tmp_path: Path) -> None: + segments = build_title_segments("Look at @nope.py here", base_dir=tmp_path) + assert segments == [TextSegment(text="Look at @nope.py here")] + + def test_folder_mention_uses_basename(self, tmp_path: Path) -> None: + folder = tmp_path / "components" + folder.mkdir() + + segments = build_title_segments("Update @components", base_dir=tmp_path) + assert segments == [ + TextSegment(text="Update "), + MentionSegment(name="components"), + ] + + def test_multiple_mentions_keep_text_in_between(self, tmp_path: Path) -> None: + (tmp_path / "a.py").write_text("", encoding="utf-8") + (tmp_path / "b.py").write_text("", encoding="utf-8") + + segments = build_title_segments("@a.py vs @b.py", base_dir=tmp_path) + assert segments == [ + MentionSegment(name="a.py"), + TextSegment(text=" vs "), + MentionSegment(name="b.py"), + ] + + def test_mention_with_no_surrounding_text(self, tmp_path: Path) -> None: + (tmp_path / "only.py").write_text("", encoding="utf-8") + + segments = build_title_segments("@only.py", base_dir=tmp_path) + assert segments == [MentionSegment(name="only.py")] diff --git a/tests/browser_sign_in/stubs.py b/tests/browser_sign_in/stubs.py index 797e1a9..177bfea 100644 --- a/tests/browser_sign_in/stubs.py +++ b/tests/browser_sign_in/stubs.py @@ -99,7 +99,7 @@ def build_sign_in_process( f"https://console.mistral.ai/codestral/cli/authenticate#{fragment}" ), poll_url=( - f"https://api.mistral.ai/api/vibe/sign-in/poll/poll-token-{process_id}" + f"https://console.mistral.ai/api/vibe/sign-in/poll/poll-token-{process_id}" ), expires_at=now + timedelta(minutes=5), ) diff --git a/tests/browser_sign_in/test_browser_sign_in.py b/tests/browser_sign_in/test_browser_sign_in.py index 55606ce..01a6ca9 100644 --- a/tests/browser_sign_in/test_browser_sign_in.py +++ b/tests/browser_sign_in/test_browser_sign_in.py @@ -34,7 +34,7 @@ TEST_SIGN_IN_URL = "https://console.mistral.ai/codestral/cli/authenticate#" + ur "complete_token": f"complete-token-{TEST_PROCESS_ID}", "state": f"state-{TEST_PROCESS_ID}", }) -TEST_POLL_URL = "https://api.mistral.ai/api/vibe/sign-in/poll/poll-token-process-1" +TEST_POLL_URL = "https://console.mistral.ai/api/vibe/sign-in/poll/poll-token-process-1" def build_code_challenge(verifier: str) -> str: @@ -94,6 +94,50 @@ async def test_authenticate_returns_api_key_after_pending_poll() -> None: assert gateway.exchange_requests[0].exchange_token == "exchange-1" +@pytest.mark.asyncio +async def test_start_attempt_returns_attempt_without_opening_browser() -> None: + opened_urls: list[str] = [] + gateway, service = build_test_service( + poll_results=[], open_browser=lambda url: opened_urls.append(url) or True + ) + + attempt = await service.start_attempt() + + assert opened_urls == [] + assert attempt.process_id == TEST_PROCESS_ID + assert attempt.sign_in_url == TEST_SIGN_IN_URL + assert attempt.poll_url == TEST_POLL_URL + assert attempt.expires_at == build_sign_in_process(TEST_NOW).expires_at + assert gateway.exchange_requests == [] + assert gateway.code_challenges == [build_code_challenge(attempt.code_verifier)] + + +@pytest.mark.asyncio +async def test_complete_attempt_returns_api_key_without_opening_browser() -> None: + opened_urls: list[str] = [] + statuses: list[BrowserSignInStatus] = [] + gateway, service = build_test_service( + poll_results=[ + BrowserSignInPollResult(status="pending"), + BrowserSignInPollResult(status="completed", exchange_token="exchange-1"), + ], + open_browser=lambda url: opened_urls.append(url) or True, + ) + attempt = await service.start_attempt() + + api_key = await service.complete_attempt(attempt, status_callback=statuses.append) + + assert api_key == "sk-browser-key" + assert opened_urls == [] + assert statuses == [ + BrowserSignInStatus.WAITING_FOR_BROWSER_SIGN_IN, + BrowserSignInStatus.EXCHANGING, + BrowserSignInStatus.COMPLETED, + ] + assert gateway.polled_urls == [TEST_POLL_URL, TEST_POLL_URL] + assert gateway.exchange_requests[0].exchange_token == "exchange-1" + + @pytest.mark.asyncio async def test_authenticate_raises_when_polling_expires() -> None: opened_urls: list[str] = [] diff --git a/tests/browser_sign_in/test_browser_sign_in_http.py b/tests/browser_sign_in/test_browser_sign_in_http.py index 9abe907..87956c3 100644 --- a/tests/browser_sign_in/test_browser_sign_in_http.py +++ b/tests/browser_sign_in/test_browser_sign_in_http.py @@ -15,9 +15,9 @@ from vibe.setup.auth import ( HttpBrowserSignInGateway, ) -AUTH_ORIGIN = "https://api.mistral.ai" +AUTH_ORIGIN = "https://console.mistral.ai" AUTH_BROWSER_BASE_URL = "https://console.mistral.ai" -AUTH_API_BASE_URL = AUTH_ORIGIN +AUTH_API_BASE_URL = "https://console.mistral.ai/api" TEST_PROCESS_ID = "process-1" TEST_COMPLETE_TOKEN = "complete-token-1" TEST_STATE = "state-1" @@ -46,7 +46,7 @@ def build_sign_in_url( def build_poll_url( *, poll_token: str = TEST_POLL_TOKEN, api_base_url: str = AUTH_API_BASE_URL ) -> str: - return f"{api_base_url}/api/vibe/sign-in/poll/{poll_token}" + return f"{api_base_url}/vibe/sign-in/poll/{poll_token}" @asynccontextmanager @@ -72,7 +72,7 @@ async def test_http_api_creates_process_with_pkce_payload() -> None: def handler(request: httpx.Request) -> httpx.Response: nonlocal captured_body - assert request.url.path == "/vibe/sign-in" + assert request.url.path == "/api/vibe/sign-in" captured_body = request.content.decode("utf-8") return httpx.Response( 200, @@ -129,7 +129,7 @@ async def test_http_api_exchanges_token_for_api_key() -> None: def handler(request: httpx.Request) -> httpx.Response: nonlocal captured_body - assert request.url.path == "/vibe/sign-in/process-1/exchange" + assert request.url.path == "/api/vibe/sign-in/process-1/exchange" captured_body = request.content.decode("utf-8") return httpx.Response(200, json={"api_key": "sk-browser-key"}) @@ -142,6 +142,28 @@ async def test_http_api_exchanges_token_for_api_key() -> None: assert '"code_verifier":"verifier-1"' in captured_body +@pytest.mark.asyncio +async def test_http_api_logs_exchange_failure_status_and_detail_without_secrets( + caplog: pytest.LogCaptureFixture, +) -> None: + exchange_token = "exchange-token-secret" + code_verifier = "verifier-secret" + + def handler(request: httpx.Request) -> httpx.Response: + assert request.url.path == "/api/vibe/sign-in/process-1/exchange" + return httpx.Response(400, json={"detail": "Invalid exchange token."}) + + async with build_gateway(handler) as gateway: + with caplog.at_level(logging.WARNING, logger="vibe"): + with pytest.raises(BrowserSignInError, match="exchange browser sign-in"): + await gateway.exchange("process-1", exchange_token, code_verifier) + + assert "status_code=400" in caplog.text + assert "response_detail=Invalid exchange token." in caplog.text + assert exchange_token not in caplog.text + assert code_verifier not in caplog.text + + @pytest.mark.asyncio async def test_http_api_translates_transport_errors() -> None: def handler(request: httpx.Request) -> httpx.Response: @@ -232,31 +254,35 @@ async def test_http_api_accepts_poll_url_under_configured_api_base_url() -> None @pytest.mark.asyncio -async def test_http_api_accepts_poll_url_under_configured_api_base_path() -> None: +async def test_http_api_accepts_backend_poll_url_under_custom_configured_api_base_path() -> ( + None +): now = datetime(2026, 3, 16, tzinfo=UTC) + browser_base_url = "https://browser-auth.example/sign-in" + api_base_url = "https://browser-auth.example/custom/api" + poll_url = f"{api_base_url}/backend-owned-poll/{TEST_POLL_TOKEN}" def handler(request: httpx.Request) -> httpx.Response: - assert request.url.path == "/v1/vibe/sign-in" + assert request.url.path == "/custom/api/vibe/sign-in" return httpx.Response( 200, json={ "process_id": TEST_PROCESS_ID, - "sign_in_url": build_sign_in_url(), - "poll_url": build_poll_url( - poll_token=TEST_POLL_TOKEN, api_base_url="https://api.mistral.ai/v1" - ), + "sign_in_url": build_sign_in_url(base_url=browser_base_url), + "poll_url": poll_url, "expires_at": _iso(now + timedelta(minutes=5)), }, ) async with build_gateway( handler, - origin="https://api.mistral.ai", - api_base_url="https://api.mistral.ai/v1", + origin="https://browser-auth.example", + browser_base_url=browser_base_url, + api_base_url=api_base_url, ) as gateway: process = await gateway.create_process("challenge-123") - assert process.poll_url == build_poll_url(api_base_url="https://api.mistral.ai/v1") + assert process.poll_url == poll_url @pytest.mark.asyncio @@ -273,7 +299,9 @@ async def test_http_api_accepts_same_origin_urls_with_explicit_default_https_por "sign_in_url": build_sign_in_url( base_url="https://console.mistral.ai:443" ), - "poll_url": build_poll_url(api_base_url="https://api.mistral.ai:443"), + "poll_url": build_poll_url( + api_base_url="https://console.mistral.ai:443/api" + ), "expires_at": _iso(now + timedelta(minutes=5)), }, ) @@ -284,7 +312,9 @@ async def test_http_api_accepts_same_origin_urls_with_explicit_default_https_por assert process.sign_in_url == build_sign_in_url( base_url="https://console.mistral.ai:443" ) - assert process.poll_url == build_poll_url(api_base_url="https://api.mistral.ai:443") + assert process.poll_url == build_poll_url( + api_base_url="https://console.mistral.ai:443/api" + ) @pytest.mark.asyncio @@ -299,8 +329,8 @@ async def test_http_api_accepts_poll_url_without_explicit_default_https_port_whe async with build_gateway( handler, - origin="https://api.mistral.ai:443", - api_base_url="https://api.mistral.ai:443", + origin="https://console.mistral.ai:443", + api_base_url="https://console.mistral.ai:443/api", ) as gateway: result = await gateway.poll(build_poll_url()) @@ -365,15 +395,13 @@ async def test_http_api_rejects_sign_in_url_outside_browser_base_path_after_norm "sign_in_url": build_sign_in_url( base_url="https://console.mistral.ai/v1/.." ), - "poll_url": build_poll_url(api_base_url="https://api.mistral.ai/v1"), + "poll_url": build_poll_url(), "expires_at": _iso(now + timedelta(minutes=5)), }, ) async with build_gateway( - handler, - browser_base_url="https://console.mistral.ai/v1", - api_base_url="https://api.mistral.ai/v1", + handler, browser_base_url="https://console.mistral.ai/v1" ) as gateway: with pytest.raises(BrowserSignInError, match="start browser sign-in") as err: await gateway.create_process("challenge-123") @@ -395,15 +423,13 @@ async def test_http_api_rejects_sign_in_url_with_encoded_dot_segments_outside_br "sign_in_url": build_sign_in_url( base_url="https://console.mistral.ai/v1/%2e%2e" ), - "poll_url": build_poll_url(api_base_url="https://api.mistral.ai/v1"), + "poll_url": build_poll_url(), "expires_at": _iso(now + timedelta(minutes=5)), }, ) async with build_gateway( - handler, - browser_base_url="https://console.mistral.ai/v1", - api_base_url="https://api.mistral.ai/v1", + handler, browser_base_url="https://console.mistral.ai/v1" ) as gateway: with pytest.raises(BrowserSignInError, match="start browser sign-in") as err: await gateway.create_process("challenge-123") @@ -449,7 +475,7 @@ async def test_http_api_rejects_returned_poll_url_outside_api_base_path_after_no ), "poll_url": build_poll_url( poll_token=TEST_POLL_TOKEN, - api_base_url="https://api.mistral.ai/v1/..", + api_base_url="https://console.mistral.ai/api/..", ), "expires_at": _iso(now + timedelta(minutes=5)), }, @@ -457,8 +483,7 @@ async def test_http_api_rejects_returned_poll_url_outside_api_base_path_after_no async with build_gateway( handler, - origin="https://api.mistral.ai", - api_base_url="https://api.mistral.ai/v1", + origin="https://console.mistral.ai", browser_base_url="https://console.mistral.ai/v1", ) as gateway: with pytest.raises(BrowserSignInError, match="start browser sign-in") as err: @@ -483,7 +508,7 @@ async def test_http_api_rejects_returned_poll_url_with_encoded_dot_segments_outs ), "poll_url": build_poll_url( poll_token=TEST_POLL_TOKEN, - api_base_url="https://api.mistral.ai/v1/%2e%2e", + api_base_url="https://console.mistral.ai/api/%2e%2e", ), "expires_at": _iso(now + timedelta(minutes=5)), }, @@ -491,8 +516,7 @@ async def test_http_api_rejects_returned_poll_url_with_encoded_dot_segments_outs async with build_gateway( handler, - origin="https://api.mistral.ai", - api_base_url="https://api.mistral.ai/v1", + origin="https://console.mistral.ai", browser_base_url="https://console.mistral.ai/v1", ) as gateway: with pytest.raises(BrowserSignInError, match="start browser sign-in") as err: @@ -504,15 +528,14 @@ async def test_http_api_rejects_returned_poll_url_with_encoded_dot_segments_outs @pytest.mark.asyncio async def test_http_api_rejects_poll_url_outside_api_base_path() -> None: async with build_gateway( - lambda _: httpx.Response(200), - origin="https://api.mistral.ai", - api_base_url="https://api.mistral.ai/v1", + lambda _: httpx.Response(200, json={"status": "completed"}), + origin="https://console.mistral.ai", ) as gateway: with pytest.raises( BrowserSignInError, match="status could not be retrieved" ) as err: await gateway.poll( - "https://api.mistral.ai/v1evil/api/vibe/sign-in/poll/poll-token-1" + "https://console.mistral.ai/apievil/vibe/sign-in/poll/poll-token-1" ) assert err.value.code is BrowserSignInErrorCode.POLL_FAILED @@ -523,15 +546,14 @@ async def test_http_api_rejects_poll_url_outside_api_base_path_after_normalizati None ): async with build_gateway( - lambda _: httpx.Response(200), - origin="https://api.mistral.ai", - api_base_url="https://api.mistral.ai/v1", + lambda _: httpx.Response(200, json={"status": "completed"}), + origin="https://console.mistral.ai", ) as gateway: with pytest.raises( BrowserSignInError, match="status could not be retrieved" ) as err: await gateway.poll( - "https://api.mistral.ai/v1/../api/vibe/sign-in/poll/poll-token-1" + "https://console.mistral.ai/api/../evil/sign-in/poll/poll-token-1" ) assert err.value.code is BrowserSignInErrorCode.POLL_FAILED @@ -547,7 +569,7 @@ async def test_http_api_translates_invalid_returned_poll_url_port() -> None: json={ "process_id": TEST_PROCESS_ID, "sign_in_url": build_sign_in_url(), - "poll_url": "https://api.mistral.ai:99999/api/vibe/sign-in/poll/poll-token-1", + "poll_url": "https://console.mistral.ai:99999/api/vibe/sign-in/poll/poll-token-1", "expires_at": _iso(now + timedelta(minutes=5)), }, ) @@ -566,7 +588,7 @@ async def test_http_api_assigns_poll_failed_code_on_invalid_poll_url_port() -> N BrowserSignInError, match="status could not be retrieved" ) as err: await gateway.poll( - "https://api.mistral.ai:99999/api/vibe/sign-in/poll/poll-token-1" + "https://console.mistral.ai:99999/api/vibe/sign-in/poll/poll-token-1" ) assert err.value.code is BrowserSignInErrorCode.POLL_FAILED @@ -639,16 +661,15 @@ async def test_http_api_does_not_log_poll_secret_on_poll_url_validation_failure( poll_token = "poll-token-secret" async with build_gateway( - lambda _: httpx.Response(200), - origin="https://api.mistral.ai", - api_base_url="https://api.mistral.ai/v1", + lambda _: httpx.Response(200, json={"status": "completed"}), + origin="https://console.mistral.ai", ) as gateway: with caplog.at_level(logging.WARNING, logger="vibe"): with pytest.raises( BrowserSignInError, match="status could not be retrieved" ): await gateway.poll( - f"https://api.mistral.ai/v1evil/api/vibe/sign-in/poll/{poll_token}" + f"https://console.mistral.ai/apievil/sign-in/poll/{poll_token}" ) assert poll_token not in caplog.text diff --git a/tests/core/session/test_title_format.py b/tests/core/session/test_title_format.py new file mode 100644 index 0000000..caf399e --- /dev/null +++ b/tests/core/session/test_title_format.py @@ -0,0 +1,74 @@ +from __future__ import annotations + +from vibe.core.session.title_format import ( + MAX_TITLE_LENGTH, + MentionSegment, + TextSegment, + format_session_title, +) + + +class TestFormatSessionTitle: + def test_empty_returns_empty(self) -> None: + assert format_session_title([]) == "" + + def test_only_whitespace_returns_empty(self) -> None: + assert format_session_title([TextSegment(text=" \n\t ")]) == "" + + def test_plain_text_passthrough(self) -> None: + assert ( + format_session_title([TextSegment(text="Refactor the auth")]) + == "Refactor the auth" + ) + + def test_mention_no_lines(self) -> None: + assert format_session_title([MentionSegment(name="foo.py")]) == "@foo.py" + + def test_mention_with_start_line_only(self) -> None: + assert ( + format_session_title([MentionSegment(name="foo.py", start_line=12)]) + == "@foo.py:12" + ) + + def test_mention_with_line_range(self) -> None: + assert ( + format_session_title([ + MentionSegment(name="foo.py", start_line=9, end_line=27) + ]) + == "@foo.py:9-27" + ) + + def test_mixed_text_and_mention(self) -> None: + segments = [ + TextSegment(text="Refactor "), + MentionSegment(name="auth.py"), + TextSegment(text=" please"), + ] + assert format_session_title(segments) == "Refactor @auth.py please" + + def test_collapses_whitespace(self) -> None: + segments = [TextSegment(text="line one\n\nline two\t\nline three")] + assert format_session_title(segments) == "line one line two line three" + + def test_strips_outer_whitespace(self) -> None: + assert ( + format_session_title([TextSegment(text=" hello world ")]) == "hello world" + ) + + def test_truncates_beyond_max(self) -> None: + long_text = "a" * (MAX_TITLE_LENGTH + 20) + result = format_session_title([TextSegment(text=long_text)]) + assert result == "a" * MAX_TITLE_LENGTH + "…" + + def test_no_truncation_at_or_below_max(self) -> None: + text = "a" * MAX_TITLE_LENGTH + assert format_session_title([TextSegment(text=text)]) == text + + def test_truncation_can_cut_inside_mention(self) -> None: + segments = [ + TextSegment(text="x" * (MAX_TITLE_LENGTH - 3)), + MentionSegment(name="long-filename.py", start_line=5, end_line=10), + ] + result = format_session_title(segments) + assert len(result) == MAX_TITLE_LENGTH + 1 + assert result.endswith("…") diff --git a/tests/core/test_config_builder.py b/tests/core/test_config_builder.py new file mode 100644 index 0000000..12fb363 --- /dev/null +++ b/tests/core/test_config_builder.py @@ -0,0 +1,322 @@ +from __future__ import annotations + +from typing import Annotated, Any + +from pydantic import Field, ValidationError +import pytest + +from vibe.core.config.builder import ConfigBuilder +from vibe.core.config.layer import ConfigLayer, RawConfig +from vibe.core.config.schema import ( + ConfigFragment, + ConfigSchema, + WithConcatMerge, + WithConflictMerge, + WithReplaceMerge, + WithShallowMerge, + WithUnionMerge, +) +from vibe.core.utils.merge import MergeConflictError + + +class FakeLayer(ConfigLayer[RawConfig]): + def __init__(self, *, name: str, data: dict[str, Any]) -> None: + super().__init__(name=name) + self._data = data + + async def _check_trust(self) -> bool: + return True + + async def _read_config(self) -> dict[str, Any]: + return dict(self._data) + + +class UntrustedFakeLayer(FakeLayer): + async def _check_trust(self) -> bool: + return False + + +class InnerFragment(ConfigFragment): + value: Annotated[str, WithReplaceMerge()] = "default" + items: Annotated[list[str], WithConcatMerge()] = Field(default_factory=list) + + +class SampleSchema(ConfigSchema): + name: Annotated[str, WithReplaceMerge()] = "unnamed" + tags: Annotated[list[str], WithConcatMerge()] = Field(default_factory=list) + entries: Annotated[list[dict[str, str]], WithUnionMerge(merge_key="id")] = Field( + default_factory=list + ) + inner: InnerFragment = Field(default_factory=InnerFragment) + + +@pytest.mark.asyncio +async def test_replace_strategy_higher_layer_wins() -> None: + builder = ConfigBuilder(SampleSchema) + builder.add_layer(FakeLayer(name="low", data={"name": "low-name"})) + builder.add_layer(FakeLayer(name="high", data={"name": "high-name"})) + config = await builder.build() + assert config.name == "high-name" + + +@pytest.mark.asyncio +async def test_concat_strategy_appends_lists() -> None: + builder = ConfigBuilder(SampleSchema) + builder.add_layer(FakeLayer(name="base", data={"tags": ["a", "b"]})) + builder.add_layer(FakeLayer(name="extra", data={"tags": ["c"]})) + config = await builder.build() + assert config.tags == ["a", "b", "c"] + + +@pytest.mark.asyncio +async def test_union_strategy_merges_by_key() -> None: + builder = ConfigBuilder(SampleSchema) + builder.add_layer( + FakeLayer( + name="base", + data={"entries": [{"id": "1", "v": "old"}, {"id": "2", "v": "keep"}]}, + ) + ) + builder.add_layer( + FakeLayer(name="override", data={"entries": [{"id": "1", "v": "new"}]}) + ) + config = await builder.build() + assert config.entries == [{"id": "1", "v": "new"}, {"id": "2", "v": "keep"}] + + +@pytest.mark.asyncio +async def test_fragment_recursion() -> None: + builder = ConfigBuilder(SampleSchema) + builder.add_layer( + FakeLayer( + name="layer1", data={"inner": {"value": "from-layer1", "items": ["x"]}} + ) + ) + builder.add_layer(FakeLayer(name="layer2", data={"inner": {"items": ["y"]}})) + config = await builder.build() + assert config.inner.value == "from-layer1" + assert config.inner.items == ["x", "y"] + + +@pytest.mark.asyncio +async def test_untrusted_layer_skipped() -> None: + builder = ConfigBuilder(SampleSchema) + builder.add_layer(FakeLayer(name="trusted", data={"name": "good"})) + builder.add_layer(UntrustedFakeLayer(name="untrusted", data={"name": "bad"})) + config = await builder.build() + assert config.name == "good" + + +@pytest.mark.asyncio +async def test_empty_layers_uses_schema_defaults() -> None: + builder = ConfigBuilder(SampleSchema) + config = await builder.build() + assert config.name == "unnamed" + assert config.tags == [] + assert config.inner.value == "default" + + +@pytest.mark.asyncio +async def test_single_layer_partial_data() -> None: + builder = ConfigBuilder(SampleSchema) + builder.add_layer(FakeLayer(name="partial", data={"name": "custom"})) + config = await builder.build() + assert config.name == "custom" + assert config.tags == [] + assert config.inner.value == "default" + + +# --- MERGE (shallow dict merge) --- + + +class MergeSchema(ConfigSchema): + settings: Annotated[dict[str, Any], WithShallowMerge()] = Field( + default_factory=dict + ) + + +@pytest.mark.asyncio +async def test_shallow_merge_combines_dicts() -> None: + builder = ConfigBuilder(MergeSchema) + builder.add_layer(FakeLayer(name="base", data={"settings": {"a": 1, "b": 2}})) + builder.add_layer(FakeLayer(name="override", data={"settings": {"b": 99, "c": 3}})) + config = await builder.build() + assert config.settings == {"a": 1, "b": 99, "c": 3} + + +@pytest.mark.asyncio +async def test_shallow_merge_single_layer() -> None: + builder = ConfigBuilder(MergeSchema) + builder.add_layer(FakeLayer(name="only", data={"settings": {"x": 1}})) + config = await builder.build() + assert config.settings == {"x": 1} + + +# --- CONFLICT --- + + +class ConflictSchema(ConfigSchema): + unique_id: Annotated[str, WithConflictMerge()] = "" + + +@pytest.mark.asyncio +async def test_conflict_single_layer_succeeds() -> None: + builder = ConfigBuilder(ConflictSchema) + builder.add_layer(FakeLayer(name="only", data={"unique_id": "abc"})) + config = await builder.build() + assert config.unique_id == "abc" + + +@pytest.mark.asyncio +async def test_conflict_two_layers_raises() -> None: + builder = ConfigBuilder(ConflictSchema) + builder.add_layer(FakeLayer(name="first", data={"unique_id": "abc"})) + builder.add_layer(FakeLayer(name="second", data={"unique_id": "def"})) + with pytest.raises(MergeConflictError): + await builder.build() + + +# --- 4-layer merge with mixed strategies --- + + +class FourLayerSchema(ConfigSchema): + active_model: Annotated[str, WithReplaceMerge()] = "default" + tags: Annotated[list[str], WithConcatMerge()] = Field(default_factory=list) + models: Annotated[list[dict[str, str]], WithUnionMerge(merge_key="alias")] = Field( + default_factory=list + ) + settings: Annotated[dict[str, Any], WithShallowMerge()] = Field( + default_factory=dict + ) + + +@pytest.mark.asyncio +async def test_four_layers_with_mixed_strategies() -> None: + builder = ConfigBuilder(FourLayerSchema) + builder.add_layer( + FakeLayer( + name="defaults", + data={ + "active_model": "model-a", + "tags": ["core"], + "models": [{"alias": "m1", "provider": "p1"}], + "settings": {"timeout": 30}, + }, + ) + ) + builder.add_layer( + FakeLayer( + name="user", + data={ + "active_model": "model-b", + "tags": ["user"], + "models": [{"alias": "m2", "provider": "p2"}], + "settings": {"timeout": 60, "debug": True}, + }, + ) + ) + builder.add_layer( + FakeLayer( + name="project", + data={ + "tags": ["project"], + "models": [{"alias": "m1", "provider": "p1-override"}], + }, + ) + ) + builder.add_layer(FakeLayer(name="cli", data={"active_model": "model-d"})) + + config = await builder.build() + + assert config.active_model == "model-d" + + assert config.tags == ["core", "user", "project"] + + assert config.models == [ + {"alias": "m1", "provider": "p1-override"}, + {"alias": "m2", "provider": "p2"}, + ] + + assert config.settings == {"timeout": 60, "debug": True} + + +# --- Validation errors --- + + +class StrictSchema(ConfigSchema): + count: Annotated[int, WithReplaceMerge()] + name: Annotated[str, WithReplaceMerge()] = "default" + + +@pytest.mark.asyncio +async def test_validation_error_on_missing_required_field() -> None: + builder = ConfigBuilder(StrictSchema) + builder.add_layer(FakeLayer(name="incomplete", data={"name": "hello"})) + with pytest.raises(ValidationError, match="count"): + await builder.build() + + +@pytest.mark.asyncio +async def test_validation_error_on_wrong_type() -> None: + builder = ConfigBuilder(StrictSchema) + builder.add_layer( + FakeLayer(name="bad-type", data={"count": "not-a-number", "name": "hello"}) + ) + with pytest.raises(ValidationError): + await builder.build() + + +# --- Fragment defaults --- + + +@pytest.mark.asyncio +async def test_fragment_defaults_when_no_layer_provides() -> None: + builder = ConfigBuilder(SampleSchema) + config = await builder.build() + + assert config.inner.value == "default" + assert config.inner.items == [] + + +# --- Edge cases --- + + +@pytest.mark.asyncio +async def test_concat_with_empty_list_from_one_layer() -> None: + builder = ConfigBuilder(SampleSchema) + builder.add_layer(FakeLayer(name="empty", data={"tags": []})) + builder.add_layer(FakeLayer(name="full", data={"tags": ["a"]})) + config = await builder.build() + assert config.tags == ["a"] + + +@pytest.mark.asyncio +async def test_union_with_empty_list_from_one_layer() -> None: + builder = ConfigBuilder(SampleSchema) + builder.add_layer(FakeLayer(name="empty", data={"entries": []})) + builder.add_layer(FakeLayer(name="full", data={"entries": [{"id": "1", "v": "a"}]})) + config = await builder.build() + assert config.entries == [{"id": "1", "v": "a"}] + + +@pytest.mark.asyncio +async def test_all_layers_untrusted_uses_defaults() -> None: + builder = ConfigBuilder(SampleSchema) + builder.add_layer(UntrustedFakeLayer(name="u1", data={"name": "bad1"})) + builder.add_layer(UntrustedFakeLayer(name="u2", data={"name": "bad2"})) + config = await builder.build() + assert config.name == "unnamed" + + +class NullableSchema(ConfigSchema): + value: Annotated[str | None, WithReplaceMerge()] = "default" + + +@pytest.mark.asyncio +async def test_replace_none_means_absent_so_base_wins() -> None: + builder = ConfigBuilder(NullableSchema) + builder.add_layer(FakeLayer(name="base", data={"value": "hello"})) + builder.add_layer(FakeLayer(name="nullifier", data={"value": None})) + config = await builder.build() + # None is treated as "not provided" by MergeStrategy, so base wins + assert config.value == "hello" diff --git a/tests/core/test_config_orchestrator.py b/tests/core/test_config_orchestrator.py new file mode 100644 index 0000000..d8a73ad --- /dev/null +++ b/tests/core/test_config_orchestrator.py @@ -0,0 +1,86 @@ +from __future__ import annotations + +from typing import Annotated, Any + +from pydantic import ValidationError +import pytest + +from vibe.core.config.layer import ConfigLayer, RawConfig +from vibe.core.config.orchestrator import ConfigOrchestrator +from vibe.core.config.schema import ConfigSchema, WithReplaceMerge + + +class FakeLayer(ConfigLayer[RawConfig]): + def __init__(self, *, name: str, data: dict[str, Any]) -> None: + super().__init__(name=name) + self._data = data + + async def _check_trust(self) -> bool: + return True + + async def _read_config(self) -> dict[str, Any]: + return dict(self._data) + + +class SimpleSchema(ConfigSchema): + value: Annotated[str, WithReplaceMerge()] = "default" + + +@pytest.mark.asyncio +async def test_create_builds_config() -> None: + layer = FakeLayer(name="test", data={"value": "hello"}) + orch = await ConfigOrchestrator.create(schema=SimpleSchema, layers=[layer]) + assert orch.config.value == "hello" + + +@pytest.mark.asyncio +async def test_get_layer_returns_named_layer() -> None: + layer = FakeLayer(name="my-layer", data={}) + orch = await ConfigOrchestrator.create(schema=SimpleSchema, layers=[layer]) + assert orch.get_layer("my-layer") is layer + + +@pytest.mark.asyncio +async def test_get_layer_unknown_raises() -> None: + orch = await ConfigOrchestrator.create(schema=SimpleSchema, layers=[]) + with pytest.raises(KeyError, match="unknown"): + orch.get_layer("unknown") + + +@pytest.mark.asyncio +async def test_reload_picks_up_changes() -> None: + layer = FakeLayer(name="mutable", data={"value": "original"}) + orch = await ConfigOrchestrator.create(schema=SimpleSchema, layers=[layer]) + assert orch.config.value == "original" + + layer._data = {"value": "updated"} + await orch.reload() + assert orch.config.value == "updated" + + +@pytest.mark.asyncio +async def test_config_is_immutable() -> None: + layer = FakeLayer(name="test", data={"value": "hello"}) + orch = await ConfigOrchestrator.create(schema=SimpleSchema, layers=[layer]) + with pytest.raises(ValidationError, match="frozen"): + orch.config.value = "changed" # type: ignore[misc] + + +@pytest.mark.asyncio +async def test_origin_of_missing_key_returns_none() -> None: + orch = await ConfigOrchestrator.create(schema=SimpleSchema, layers=[]) + assert orch.config.origin_of("nonexistent") is None + + +@pytest.mark.asyncio +async def test_apply_patch_raises_not_implemented() -> None: + orch = await ConfigOrchestrator.create(schema=SimpleSchema, layers=[]) + with pytest.raises(NotImplementedError, match="M2"): + await orch.apply_patch({}) + + +@pytest.mark.asyncio +async def test_subscribe_raises_not_implemented() -> None: + orch = await ConfigOrchestrator.create(schema=SimpleSchema, layers=[]) + with pytest.raises(NotImplementedError, match="M3"): + await orch.subscribe(lambda: None) diff --git a/tests/core/test_config_toml_end_to_end.py b/tests/core/test_config_toml_end_to_end.py new file mode 100644 index 0000000..e1de0dc --- /dev/null +++ b/tests/core/test_config_toml_end_to_end.py @@ -0,0 +1,71 @@ +from __future__ import annotations + +from pathlib import Path +from typing import Annotated + +from pydantic import Field, ValidationError +import pytest + +from vibe.core.config.schema import ( + ConfigFragment, + ConfigSchema, + WithReplaceMerge, + WithUnionMerge, +) + + +class ModelsFragment(ConfigFragment): + active_model: Annotated[str, WithReplaceMerge()] = "default-model" + models: Annotated[list[dict[str, str]], WithUnionMerge(merge_key="alias")] = Field( + default_factory=list + ) + providers: Annotated[list[dict[str, str]], WithUnionMerge(merge_key="name")] = ( + Field(default_factory=list) + ) + + +class MinimalSchema(ConfigSchema): + models: ModelsFragment = Field(default_factory=ModelsFragment) + + +@pytest.mark.asyncio +async def test_toml_to_typed_config_end_to_end(tmp_working_directory: Path) -> None: + toml_path = tmp_working_directory / "config.toml" + toml_path.write_text( + """\ +[models] +active_model = "mistral-large" + +[[models.models]] +alias = "mistral-large" +provider = "mistral" + +[[models.providers]] +name = "mistral" +api_base = "https://api.mistral.ai/v1" +""" + ) + + from vibe.core.config.layers.user import UserConfigLayer + from vibe.core.config.orchestrator import ConfigOrchestrator + + layer = UserConfigLayer(path=toml_path, name="user-toml") + orchestrator = await ConfigOrchestrator.create(schema=MinimalSchema, layers=[layer]) + + assert orchestrator.config.models.active_model == "mistral-large" + assert orchestrator.config.models.models == [ + {"alias": "mistral-large", "provider": "mistral"} + ] + # Immutability (frozen schema) + with pytest.raises(ValidationError, match="frozen"): + orchestrator.config.models = ModelsFragment() # type: ignore[misc] + + # Reload picks up changes + toml_path.write_text( + """\ +[models] +active_model = "codestral" +""" + ) + await orchestrator.reload() + assert orchestrator.config.models.active_model == "codestral" diff --git a/tests/core/test_user_config_layer.py b/tests/core/test_user_config_layer.py new file mode 100644 index 0000000..755f9f4 --- /dev/null +++ b/tests/core/test_user_config_layer.py @@ -0,0 +1,96 @@ +from __future__ import annotations + +from pathlib import Path + +import pytest + +from vibe.core.config.layer import LayerImplementationError +from vibe.core.config.layers.user import UserConfigLayer + + +@pytest.mark.asyncio +async def test_reads_toml_file(tmp_working_directory: Path) -> None: + path = tmp_working_directory / "config.toml" + path.write_text('active_model = "mistral-large"\ncount = 42\n') + + layer = UserConfigLayer(path=path, name="user-toml") + data = await layer.load() + assert data.model_extra == {"active_model": "mistral-large", "count": 42} + + +@pytest.mark.asyncio +async def test_always_trusted(tmp_working_directory: Path) -> None: + path = tmp_working_directory / "config.toml" + path.write_text('key = "value"\n') + + layer = UserConfigLayer(path=path, name="user-toml") + assert layer.is_trusted is None + data = await layer.load() + assert layer.is_trusted is True + assert data.model_extra == {"key": "value"} + + +@pytest.mark.asyncio +async def test_missing_file_returns_empty(tmp_working_directory: Path) -> None: + path = tmp_working_directory / "nonexistent.toml" + layer = UserConfigLayer(path=path, name="user-toml") + data = await layer.load() + assert data.model_extra == {} + + +@pytest.mark.asyncio +async def test_apply_raises_not_implemented(tmp_working_directory: Path) -> None: + path = tmp_working_directory / "config.toml" + layer = UserConfigLayer(path=path, name="user-toml") + with pytest.raises(NotImplementedError, match="M2"): + await layer.apply({"op": "set"}) + + +@pytest.mark.asyncio +async def test_nested_toml_structure(tmp_working_directory: Path) -> None: + path = tmp_working_directory / "config.toml" + path.write_text("""\ +[models] +active_model = "test" + +[[models.items]] +alias = "a" +provider = "p" +""") + layer = UserConfigLayer(path=path, name="user-toml") + data = await layer.load() + assert data.model_extra == { + "models": {"active_model": "test", "items": [{"alias": "a", "provider": "p"}]} + } + + +@pytest.mark.asyncio +async def test_invalid_toml_raises(tmp_working_directory: Path) -> None: + path = tmp_working_directory / "bad.toml" + path.write_text("this is not valid = = = toml [[[") + layer = UserConfigLayer(path=path, name="user-toml") + with pytest.raises(LayerImplementationError, match="_read_config"): + await layer.load() + + +@pytest.mark.asyncio +async def test_force_reload_reads_fresh_data(tmp_working_directory: Path) -> None: + path = tmp_working_directory / "config.toml" + path.write_text('value = "first"\n') + layer = UserConfigLayer(path=path, name="user-toml") + + data1 = await layer.load() + assert data1.model_extra == {"value": "first"} + + path.write_text('value = "second"\n') + data2 = await layer.load(force=True) + assert data2.model_extra == {"value": "second"} + + +@pytest.mark.asyncio +async def test_empty_toml_file(tmp_working_directory: Path) -> None: + path = tmp_working_directory / "empty.toml" + path.write_text("") + layer = UserConfigLayer(path=path, name="user-toml") + data = await layer.load() + assert data.model_extra == {} diff --git a/tests/onboarding/test_ui_onboarding.py b/tests/onboarding/test_ui_onboarding.py index a27e9d1..e46b327 100644 --- a/tests/onboarding/test_ui_onboarding.py +++ b/tests/onboarding/test_ui_onboarding.py @@ -3,7 +3,6 @@ from __future__ import annotations import asyncio from collections.abc import Callable from pathlib import Path -from types import SimpleNamespace from typing import cast import pytest @@ -31,9 +30,10 @@ from vibe.setup.auth import ( BrowserSignInService, BrowserSignInStatus, ) +from vibe.setup.auth.api_key_persistence import persist_api_key import vibe.setup.onboarding as onboarding_module from vibe.setup.onboarding import OnboardingApp -from vibe.setup.onboarding.screens.api_key import ApiKeyScreen, persist_api_key +from vibe.setup.onboarding.screens.api_key import ApiKeyScreen from vibe.setup.onboarding.screens.auth_method import AuthMethodScreen from vibe.setup.onboarding.screens.browser_sign_in import ( ERROR_HINT, @@ -43,7 +43,7 @@ from vibe.setup.onboarding.screens.browser_sign_in import ( ) CONSOLE_URL = "https://console.mistral.ai" -API_URL = "https://api.mistral.ai" +BROWSER_AUTH_API_URL = "https://console.mistral.ai/api" async def _wait_for( @@ -95,7 +95,7 @@ def _build_browser_onboarding_app( return OnboardingApp( config=_build_onboarding_config( browser_auth_base_url=CONSOLE_URL, - browser_auth_api_base_url=API_URL, + browser_auth_api_base_url=BROWSER_AUTH_API_URL, enable_experimental_browser_sign_in=True, ), browser_sign_in_service_factory=browser_sign_in_service_factory, @@ -237,7 +237,8 @@ async def test_ui_hides_browser_sign_in_when_experimental_flag_is_disabled() -> ) app = OnboardingApp( config=_build_onboarding_config( - browser_auth_base_url=CONSOLE_URL, browser_auth_api_base_url=API_URL + browser_auth_base_url=CONSOLE_URL, + browser_auth_api_base_url=BROWSER_AUTH_API_URL, ), browser_sign_in_service_factory=browser_sign_in_service_factory, ) @@ -257,7 +258,7 @@ async def test_ui_supports_browser_sign_in_when_experimental_flag_is_enabled() - app = OnboardingApp( config=_build_onboarding_config( browser_auth_base_url=CONSOLE_URL, - browser_auth_api_base_url=API_URL, + browser_auth_api_base_url=BROWSER_AUTH_API_URL, enable_experimental_browser_sign_in=True, ), browser_sign_in_service_factory=browser_sign_in_service_factory, @@ -276,7 +277,7 @@ async def test_ui_offers_browser_sign_in_for_renamed_mistral_provider() -> None: provider_name="customer-mistral", backend=Backend.MISTRAL, browser_auth_base_url=CONSOLE_URL, - browser_auth_api_base_url=API_URL, + browser_auth_api_base_url=BROWSER_AUTH_API_URL, enable_experimental_browser_sign_in=True, ) ) @@ -344,7 +345,7 @@ async def test_ui_browser_sign_in_falls_back_to_mistral_env_var_when_missing() - provider_name="custom-mistral", api_key_env_var="", browser_auth_base_url=CONSOLE_URL, - browser_auth_api_base_url=API_URL, + browser_auth_api_base_url=BROWSER_AUTH_API_URL, enable_experimental_browser_sign_in=True, ), browser_sign_in_service_factory=browser_sign_in_service_factory, @@ -683,7 +684,7 @@ async def test_ui_falls_back_to_default_onboarding_context_with_invalid_active_m 'api_base = "https://api.mistral.ai/v1"', 'api_key_env_var = "MISTRAL_API_KEY"', 'browser_auth_base_url = "https://console.mistral.ai"', - 'browser_auth_api_base_url = "https://api.mistral.ai"', + 'browser_auth_api_base_url = "https://console.mistral.ai/api"', 'backend = "mistral"', "", "[[models]]", @@ -732,11 +733,9 @@ def test_api_key_screen_uses_mistral_fallback_for_context_without_env_key( monkeypatch: pytest.MonkeyPatch, ) -> None: monkeypatch.setattr( - "vibe.setup.onboarding.screens.api_key.OnboardingContext.load", - lambda: SimpleNamespace( - provider=ProviderConfig( - name="llamacpp", api_base="http://127.0.0.1:8080/v1", api_key_env_var="" - ) + "vibe.setup.auth.api_key_persistence._load_onboarding_provider", + lambda: ProviderConfig( + name="llamacpp", api_base="http://127.0.0.1:8080/v1", api_key_env_var="" ), ) @@ -778,7 +777,7 @@ def test_persist_api_key_sends_onboarding_telemetry_with_entrypoint_metadata( provider = ProviderConfig( name="mistral", - api_base="https://api.mistral.ai/v1", + api_base="https://inference.mistral.test/v1", api_key_env_var="MISTRAL_API_KEY", backend=Backend.MISTRAL, ) diff --git a/tests/session/test_session_logger.py b/tests/session/test_session_logger.py index bcf1605..18e3c76 100644 --- a/tests/session/test_session_logger.py +++ b/tests/session/test_session_logger.py @@ -204,6 +204,84 @@ class TestSessionLoggerTitleManagement: assert logger.session_metadata.end_time == "2026-01-01T10:00:00+00:00" + def test_set_initial_auto_title_applies_when_no_title_set( + self, session_config: SessionLoggingConfig + ) -> None: + logger = SessionLogger(session_config, "test-session-123") + + applied = logger.set_initial_auto_title("Pretty title") + + assert applied is True + assert logger.session_metadata is not None + assert logger.session_metadata.title == "Pretty title" + assert logger.session_metadata.title_source == "auto" + + def test_set_initial_auto_title_noop_when_title_already_set( + self, session_config: SessionLoggingConfig + ) -> None: + logger = SessionLogger(session_config, "test-session-123") + logger.set_title("Manual title") + + applied = logger.set_initial_auto_title("Pretty title") + + assert applied is False + assert logger.session_metadata is not None + assert logger.session_metadata.title == "Manual title" + assert logger.session_metadata.title_source == "manual" + + def test_set_initial_auto_title_noop_when_prior_auto_title_set( + self, session_config: SessionLoggingConfig + ) -> None: + logger = SessionLogger(session_config, "test-session-123") + logger.set_initial_auto_title("First title") + + applied = logger.set_initial_auto_title("Second title") + + assert applied is False + assert logger.session_metadata is not None + assert logger.session_metadata.title == "First title" + + def test_set_initial_auto_title_rejects_blank( + self, session_config: SessionLoggingConfig + ) -> None: + logger = SessionLogger(session_config, "test-session-123") + + applied = logger.set_initial_auto_title(" ") + + assert applied is False + assert logger.session_metadata is not None + assert logger.session_metadata.title is None + + def test_needs_initial_auto_title_true_when_no_title( + self, session_config: SessionLoggingConfig + ) -> None: + logger = SessionLogger(session_config, "test-session-123") + + assert logger.needs_initial_auto_title() is True + + def test_needs_initial_auto_title_false_after_set_initial_auto_title( + self, session_config: SessionLoggingConfig + ) -> None: + logger = SessionLogger(session_config, "test-session-123") + logger.set_initial_auto_title("Pretty title") + + assert logger.needs_initial_auto_title() is False + + def test_needs_initial_auto_title_false_after_manual_set_title( + self, session_config: SessionLoggingConfig + ) -> None: + logger = SessionLogger(session_config, "test-session-123") + logger.set_title("Manual title") + + assert logger.needs_initial_auto_title() is False + + def test_needs_initial_auto_title_false_when_disabled( + self, disabled_session_config: SessionLoggingConfig + ) -> None: + logger = SessionLogger(disabled_session_config, "test-session-123") + + assert logger.needs_initial_auto_title() is False + class TestSessionLoggerSaveInteraction: @pytest.mark.asyncio @@ -544,6 +622,47 @@ class TestSessionLoggerSaveInteraction: assert metadata["title"] == expected_title assert metadata["title_source"] == "auto" + @pytest.mark.asyncio + async def test_save_interaction_preserves_preset_auto_title( + self, + session_config: SessionLoggingConfig, + mock_vibe_config: VibeConfig, + mock_tool_manager: ToolManager, + mock_agent_profile: AgentProfile, + ) -> None: + session_id = "test-session-123" + logger = SessionLogger(session_config, session_id) + assert logger.session_metadata is not None + + logger.set_initial_auto_title("Pretty @foo.py title") + + messages = [ + LLMMessage(role=Role.system, content="System prompt"), + LLMMessage( + role=Role.user, content="path: file:///abs/foo.py\ncontent: ..." + ), + 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" + with open(metadata_file) as f: + metadata = json.load(f) + + assert metadata["title"] == "Pretty @foo.py title" + assert metadata["title_source"] == "auto" + @pytest.mark.asyncio async def test_save_interaction_preserves_manual_title( self, diff --git a/tests/test_agent_auto_title.py b/tests/test_agent_auto_title.py new file mode 100644 index 0000000..b54e34b --- /dev/null +++ b/tests/test_agent_auto_title.py @@ -0,0 +1,83 @@ +from __future__ import annotations + +from pathlib import Path + +import pytest + +from tests.conftest import build_test_agent_loop, build_test_vibe_config +from tests.mock.utils import mock_llm_chunk +from tests.stubs.fake_backend import FakeBackend +from vibe.core.agent_loop import AgentLoop +from vibe.core.config import SessionLoggingConfig +from vibe.core.types import BaseEvent, SessionTitleUpdatedEvent, UserMessageEvent + + +def _make_agent_loop(tmp_path: Path) -> AgentLoop: + session_logging = SessionLoggingConfig( + save_dir=str(tmp_path / "sessions"), session_prefix="session", enabled=True + ) + config = build_test_vibe_config(session_logging=session_logging) + backend = FakeBackend([ + [mock_llm_chunk(content="ok")], + [mock_llm_chunk(content="ok")], + ]) + return build_test_agent_loop(config=config, backend=backend) + + +async def _collect(loop: AgentLoop, prompt: str, **kwargs) -> list[BaseEvent]: + return [ev async for ev in loop.act(prompt, **kwargs)] + + +class TestAgentLoopAutoTitleEvent: + @pytest.mark.asyncio + async def test_emits_event_on_first_user_message(self, tmp_path: Path) -> None: + loop = _make_agent_loop(tmp_path) + + events = await _collect(loop, "rendered prompt", auto_title="Pretty title") + + title_events = [e for e in events if isinstance(e, SessionTitleUpdatedEvent)] + assert len(title_events) == 1 + assert title_events[0].title == "Pretty title" + + @pytest.mark.asyncio + async def test_event_fires_after_user_message_event(self, tmp_path: Path) -> None: + loop = _make_agent_loop(tmp_path) + + events = await _collect(loop, "rendered", auto_title="Pretty") + + indices = { + type(e).__name__: i + for i, e in enumerate(events) + if isinstance(e, (UserMessageEvent, SessionTitleUpdatedEvent)) + } + assert indices["UserMessageEvent"] < indices["SessionTitleUpdatedEvent"] + + @pytest.mark.asyncio + async def test_no_event_on_second_message(self, tmp_path: Path) -> None: + loop = _make_agent_loop(tmp_path) + await _collect(loop, "first", auto_title="First title") + + events = await _collect(loop, "second", auto_title="Second title") + + title_events = [e for e in events if isinstance(e, SessionTitleUpdatedEvent)] + assert title_events == [] + + @pytest.mark.asyncio + async def test_no_event_when_auto_title_is_none(self, tmp_path: Path) -> None: + loop = _make_agent_loop(tmp_path) + + events = await _collect(loop, "rendered", auto_title=None) + + title_events = [e for e in events if isinstance(e, SessionTitleUpdatedEvent)] + assert title_events == [] + + @pytest.mark.asyncio + async def test_no_event_when_session_logging_disabled(self, tmp_path: Path) -> None: + config = build_test_vibe_config() + backend = FakeBackend(mock_llm_chunk(content="ok")) + loop = build_test_agent_loop(config=config, backend=backend) + + events = await _collect(loop, "rendered", auto_title="Pretty title") + + title_events = [e for e in events if isinstance(e, SessionTitleUpdatedEvent)] + assert title_events == [] diff --git a/uv.lock b/uv.lock index f3bb81e..14a21fb 100644 --- a/uv.lock +++ b/uv.lock @@ -3,7 +3,7 @@ revision = 3 requires-python = ">=3.12" [options] -exclude-newer = "0001-01-01T00:00:00Z" # This has no effect and is included for backwards compatibility when using relative exclude-newer values. +exclude-newer = "2026-05-13T08:56:52.756107Z" exclude-newer-span = "P7D" [options.exclude-newer-package] @@ -820,7 +820,7 @@ wheels = [ [[package]] name = "mistral-vibe" -version = "2.10.0" +version = "2.10.1" source = { editable = "." } dependencies = [ { name = "agent-client-protocol" }, diff --git a/vibe/__init__.py b/vibe/__init__.py index 8139979..972866a 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.10.0" +__version__ = "2.10.1" diff --git a/vibe/acp/acp_agent_loop.py b/vibe/acp/acp_agent_loop.py index 44f22ab..2dc9c97 100644 --- a/vibe/acp/acp_agent_loop.py +++ b/vibe/acp/acp_agent_loop.py @@ -3,12 +3,14 @@ from __future__ import annotations import asyncio from collections.abc import AsyncGenerator, Callable from contextlib import aclosing +from dataclasses import dataclass +from datetime import UTC import logging import os from pathlib import Path import signal import sys -from typing import Any, cast, override +from typing import Any, Protocol, cast, override from uuid import uuid4 from acp import ( @@ -57,6 +59,7 @@ from acp.schema import ( SetSessionConfigOptionResponse, SseMcpServer, TerminalAuthMethod, + TerminalToolCallContent, TextContentBlock, TextResourceContents, ToolCallProgress, @@ -83,7 +86,9 @@ from vibe.acp.exceptions import ( UnauthenticatedError, ) from vibe.acp.session import AcpSessionLoop +from vibe.acp.title import acp_blocks_to_title_segments from vibe.acp.tools.base import BaseAcpTool +from vibe.acp.tools.events import ToolTerminalOpenedEvent from vibe.acp.tools.session_update import ( resolve_kind, tool_call_session_update, @@ -112,6 +117,7 @@ from vibe.core.agents.models import CHAT as CHAT_AGENT, BuiltinAgentName from vibe.core.autocompletion.path_prompt_adapter import render_path_prompt from vibe.core.config import ( MissingAPIKeyError, + ProviderConfig, SessionLoggingConfig, VibeConfig, load_dotenv_values, @@ -129,6 +135,7 @@ from vibe.core.session.saved_sessions import ( update_saved_session_title_at_path, ) from vibe.core.session.session_loader import SessionLoader +from vibe.core.session.title_format import format_session_title from vibe.core.skills.manager import SkillManager from vibe.core.telemetry.build_metadata import build_entrypoint_metadata from vibe.core.telemetry.send import TelemetryClient @@ -146,6 +153,7 @@ from vibe.core.types import ( RateLimitError as CoreRateLimitError, ReasoningEvent, Role, + SessionTitleUpdatedEvent, ToolCallEvent, ToolResultEvent, ToolStreamEvent, @@ -155,6 +163,18 @@ from vibe.core.utils import ( ConversationLimitException, get_user_cancellation_message, ) +from vibe.setup.auth import ( + BrowserSignInAttempt, + BrowserSignInError, + BrowserSignInErrorCode, + BrowserSignInService, + HttpBrowserSignInGateway, +) +from vibe.setup.auth.api_key_persistence import ( + persist_api_key, + resolve_api_key_provider, +) +from vibe.setup.onboarding.context import OnboardingContext logger = logging.getLogger("vibe") @@ -215,13 +235,123 @@ def _resolved_user_message_id(client_message_id: str | None) -> str: return str(uuid4()) +@dataclass(frozen=True) +class PendingBrowserSignInAttempt: + attempt: BrowserSignInAttempt + provider: ProviderConfig + + +RETRYABLE_BROWSER_SIGN_IN_COMPLETION_ERRORS = { + BrowserSignInErrorCode.EXCHANGE_FAILED, + BrowserSignInErrorCode.POLL_FAILED, +} + + +OnboardingContextLoader = Callable[[], OnboardingContext] +ApiKeyPersister = Callable[[ProviderConfig, str], str] + + +class BrowserSignInServiceAdapter(Protocol): + async def authenticate(self) -> str: ... + + async def start_attempt(self) -> BrowserSignInAttempt: ... + + async def complete_attempt(self, attempt: BrowserSignInAttempt) -> str: ... + + async def aclose(self) -> None: ... + + +BrowserSignInServiceFactory = Callable[[ProviderConfig], BrowserSignInServiceAdapter] + + class VibeAcpAgentLoop(AcpAgent): client: Client - def __init__(self) -> None: + def __init__( + self, + *, + onboarding_context_loader: OnboardingContextLoader | None = None, + browser_sign_in_service_factory: BrowserSignInServiceFactory | None = None, + api_key_persister: ApiKeyPersister = persist_api_key, + ) -> None: self.sessions: dict[str, AcpSessionLoop] = {} self.client_capabilities: ClientCapabilities | None = None self.client_info: Implementation | None = None + self._pending_browser_sign_in_attempts: dict[ + str, PendingBrowserSignInAttempt + ] = {} + self._load_onboarding_context = ( + onboarding_context_loader or OnboardingContext.load + ) + self._browser_sign_in_service_factory = ( + browser_sign_in_service_factory or self._build_browser_sign_in_service + ) + self._persist_api_key = api_key_persister + + def _build_browser_auth_method( + self, context: OnboardingContext, method_id: str + ) -> AuthMethodAgent | None: + if not context.browser_sign_in_enabled: + return None + + return AuthMethodAgent( + id=method_id, + name="Sign in through Mistral AI Studio", + description="Sign into Mistral Vibe through your Mistral AI Studio account.", + ) + + def _build_terminal_auth_method( + self, command: str, args: list[str] + ) -> TerminalAuthMethod: + return TerminalAuthMethod( + type="terminal", + id="vibe-setup", + name="Register your API Key", + description="Register your API Key inside Mistral Vibe", + field_meta={ + "terminal-auth": { + "command": command, + "args": args, + "label": "Mistral Vibe Setup", + } + }, + ) + + def _build_browser_sign_in_service( + self, provider: ProviderConfig | None = None + ) -> BrowserSignInService: + provider = provider or self._load_onboarding_context().provider + if not provider.supports_browser_sign_in: + raise InvalidRequestError( + "Browser sign-in is not available for the configured provider." + ) + + browser_base_url = provider.browser_auth_base_url + api_base_url = provider.browser_auth_api_base_url + if browser_base_url is None or api_base_url is None: + raise ConfigurationError("Browser sign-in requires both browser auth URLs.") + + return BrowserSignInService( + HttpBrowserSignInGateway( + browser_base_url=browser_base_url, api_base_url=api_base_url + ) + ) + + def _load_enabled_browser_sign_in_context(self) -> OnboardingContext: + context = self._load_onboarding_context() + if not context.browser_sign_in_enabled: + raise InvalidRequestError( + "Browser sign-in is not available for the configured provider." + ) + return context + + def _supports_delegated_browser_auth(self) -> bool: + return bool( + self.client_capabilities + and self.client_capabilities.field_meta + and self.client_capabilities.field_meta.get("browser-auth-delegated") + is True + ) @override async def initialize( @@ -256,25 +386,21 @@ class VibeAcpAgentLoop(AcpAgent): and self.client_capabilities.field_meta.get("terminal-auth") is True ) - auth_methods: list[EnvVarAuthMethod | TerminalAuthMethod | AuthMethodAgent] = ( - [ - TerminalAuthMethod( - type="terminal", - id="vibe-setup", - name="Register your API Key", - description="Register your API Key inside Mistral Vibe", - field_meta={ - "terminal-auth": { - "command": command, - "args": args, - "label": "Mistral Vibe Setup", - } - }, - ) - ] - if supports_terminal_auth - else [] - ) + context = self._load_onboarding_context() + + auth_methods: list[EnvVarAuthMethod | TerminalAuthMethod | AuthMethodAgent] = [] + if browser_auth_method := self._build_browser_auth_method( + context, "browser-auth" + ): + auth_methods.append(browser_auth_method) + if self._supports_delegated_browser_auth(): + delegated_browser_auth_method = self._build_browser_auth_method( + context, "browser-auth-delegated" + ) + if delegated_browser_auth_method is not None: + auth_methods.append(delegated_browser_auth_method) + if supports_terminal_auth: + auth_methods.append(self._build_terminal_auth_method(command, args)) response = InitializeResponse( agent_capabilities=AgentCapabilities( @@ -298,11 +424,116 @@ class VibeAcpAgentLoop(AcpAgent): ) return response + async def _authenticate_browser_auth(self, **kwargs: Any) -> AuthenticateResponse: + action = kwargs.get("action") + if action not in {None, "start"}: + raise InvalidRequestError(f"Unsupported browser auth action: {action}") + + provider = self._load_enabled_browser_sign_in_context().provider + browser_sign_in = self._browser_sign_in_service_factory(provider) + try: + api_key = await browser_sign_in.authenticate() + except BrowserSignInError as e: + raise InternalError(str(e)) from e + finally: + await browser_sign_in.aclose() + + persist_result = self._persist_api_key( + resolve_api_key_provider(provider), api_key + ) + return AuthenticateResponse( + field_meta={ + "browser-auth": {"persistResult": persist_result, "status": "completed"} + } + ) + + async def _start_delegated_browser_auth(self) -> AuthenticateResponse: + provider = self._load_enabled_browser_sign_in_context().provider + browser_sign_in = self._browser_sign_in_service_factory(provider) + try: + attempt = await browser_sign_in.start_attempt() + except BrowserSignInError as e: + raise InternalError(str(e)) from e + finally: + await browser_sign_in.aclose() + + self._pending_browser_sign_in_attempts[attempt.process_id] = ( + PendingBrowserSignInAttempt(attempt=attempt, provider=provider) + ) + return AuthenticateResponse( + field_meta={ + "browser-auth-delegated": { + "attemptId": attempt.process_id, + "expiresAt": attempt.expires_at.astimezone(UTC) + .isoformat() + .replace("+00:00", "Z"), + "signInUrl": attempt.sign_in_url, + } + } + ) + + async def _complete_delegated_browser_auth( + self, **kwargs: Any + ) -> AuthenticateResponse: + attempt_id = kwargs.get("attemptId") or kwargs.get("attempt_id") + if not isinstance(attempt_id, str) or not attempt_id: + raise InvalidRequestError("Missing browser sign-in attempt ID.") + + pending_attempt = self._pending_browser_sign_in_attempts.get(attempt_id) + if pending_attempt is None: + raise InvalidRequestError(f"Unknown browser sign-in attempt: {attempt_id}") + + browser_sign_in = self._browser_sign_in_service_factory( + pending_attempt.provider + ) + try: + api_key = await browser_sign_in.complete_attempt(pending_attempt.attempt) + except BrowserSignInError as e: + if e.code not in RETRYABLE_BROWSER_SIGN_IN_COMPLETION_ERRORS: + self._pending_browser_sign_in_attempts.pop(attempt_id, None) + raise InvalidRequestError(str(e)) from e + finally: + await browser_sign_in.aclose() + + self._pending_browser_sign_in_attempts.pop(attempt_id, None) + persist_result = self._persist_api_key( + resolve_api_key_provider(pending_attempt.provider), api_key + ) + return AuthenticateResponse( + field_meta={ + "browser-auth-delegated": { + "attemptId": attempt_id, + "persistResult": persist_result, + "status": "completed", + } + } + ) + + async def _authenticate_delegated_browser_auth( + self, **kwargs: Any + ) -> AuthenticateResponse: + action = kwargs.get("action", "start") + if action not in {"start", "complete"}: + raise InvalidRequestError( + f"Unsupported delegated browser auth action: {action}" + ) + + if action == "start": + return await self._start_delegated_browser_auth() + + return await self._complete_delegated_browser_auth(**kwargs) + @override async def authenticate( self, method_id: str, **kwargs: Any ) -> AuthenticateResponse | None: - raise NotImplementedMethodError("authenticate") + if method_id == "browser-auth": + return await self._authenticate_browser_auth(**kwargs) + + if method_id == "browser-auth-delegated": + return await self._authenticate_delegated_browser_auth(**kwargs) + + raise InvalidRequestError(f"Unsupported auth method: {method_id}") def _build_entrypoint_metadata(self) -> EntrypointMetadata: return build_entrypoint_metadata( @@ -834,9 +1065,15 @@ class VibeAcpAgentLoop(AcpAgent): ) text_prompt = SkillManager.build_skill_prompt(text_prompt, skill) + auto_title: str | None = None + if session.agent_loop.session_logger.needs_initial_auto_title(): + auto_title = ( + format_session_title(acp_blocks_to_title_segments(prompt)) or None + ) + async def agent_loop_task() -> None: async for update in self._run_agent_loop( - session, text_prompt, resolved_message_id + session, text_prompt, resolved_message_id, auto_title=auto_title ): await self.client.session_update(session_id=session.id, update=update) @@ -945,15 +1182,29 @@ class VibeAcpAgentLoop(AcpAgent): return await handler(session, text_prompt, message_id) async def _run_agent_loop( - self, session: AcpSessionLoop, prompt: str, client_message_id: str | None = None + self, + session: AcpSessionLoop, + prompt: str, + client_message_id: str | None = None, + *, + auto_title: str | None = None, ) -> AsyncGenerator[SessionUpdate | UsageUpdate]: rendered_prompt = render_path_prompt(prompt, base_dir=Path.cwd()) async with aclosing( - session.agent_loop.act(rendered_prompt, client_message_id=client_message_id) + session.agent_loop.act( + rendered_prompt, + client_message_id=client_message_id, + auto_title=auto_title, + ) ) as events: async for event in events: - if isinstance(event, AssistantEvent): + if isinstance(event, SessionTitleUpdatedEvent): + await self._emit_session_info_update( + session.id, title=event.title, updated_at=None + ) + + elif isinstance(event, AssistantEvent): yield AgentMessageChunk( session_update="agent_message_chunk", content=TextContentBlock(type="text", text=event.content), @@ -973,7 +1224,6 @@ class VibeAcpAgentLoop(AcpAgent): tool_manager=session.agent_loop.tool_manager, client=self.client, session_id=session.id, - tool_call_id=event.tool_call_id, ) session_update = tool_call_session_update(event) @@ -986,6 +1236,22 @@ class VibeAcpAgentLoop(AcpAgent): yield session_update self._send_usage_update(session) + elif isinstance(event, ToolTerminalOpenedEvent): + # bash yielded the terminal id; surface it as an + # in_progress update carrying the terminal block. + yield ToolCallProgress( + session_update="tool_call_update", + tool_call_id=event.tool_call_id, + status="in_progress", + kind=resolve_kind(event.tool_name), + content=[ + TerminalToolCallContent( + type="terminal", terminal_id=event.terminal_id + ) + ], + field_meta={"tool_name": event.tool_name}, + ) + elif isinstance(event, ToolStreamEvent): yield ToolCallProgress( session_update="tool_call_update", diff --git a/vibe/acp/title.py b/vibe/acp/title.py new file mode 100644 index 0000000..1fd2bc8 --- /dev/null +++ b/vibe/acp/title.py @@ -0,0 +1,69 @@ +from __future__ import annotations + +from pathlib import Path +import re +from urllib.parse import SplitResult, urlsplit + +from acp.helpers import ContentBlock + +from vibe.core.session.title_format import MentionSegment, TextSegment, TitleSegment + +_LINE_FRAGMENT_RE = re.compile(r"^L(\d+)(?:-L(\d+))?$") + + +def acp_blocks_to_title_segments(blocks: list[ContentBlock]) -> list[TitleSegment]: + segments: list[TitleSegment] = [] + for block in blocks: + if block.field_meta and block.field_meta.get("automatic"): + continue + if (segment := _block_to_segment(block)) is not None: + segments.append(segment) + return segments + + +def _block_to_segment(block: ContentBlock) -> TitleSegment | None: + match block.type: + case "text": + return TextSegment(text=block.text) + case "resource": + base, start, end = _parse_line_range_fragment(block.resource.uri) + name = _basename_from_uri(base) + if not name: + return None + return MentionSegment(name=name, start_line=start, end_line=end) + case "resource_link": + base, start, end = _parse_line_range_fragment(block.uri) + name = block.name or _basename_from_uri(base) + if not name: + return None + return MentionSegment(name=name, start_line=start, end_line=end) + case _: + return None + + +def _safe_urlsplit(uri: str) -> SplitResult | None: + try: + return urlsplit(uri) + except ValueError: + return None + + +def _parse_line_range_fragment(uri: str) -> tuple[str, int | None, int | None]: + parts = _safe_urlsplit(uri) + if parts is None or not parts.fragment: + return uri, None, None + match = _LINE_FRAGMENT_RE.match(parts.fragment) + if match is None: + return uri, None, None + start = int(match.group(1)) + end = int(match.group(2)) if match.group(2) is not None else None + base_uri = parts._replace(fragment="").geturl() + return base_uri, start, end + + +def _basename_from_uri(uri: str) -> str: + parts = _safe_urlsplit(uri) + if parts is None: + return "" + path = parts.path if parts.scheme else uri + return Path(path).name diff --git a/vibe/acp/tools/base.py b/vibe/acp/tools/base.py index 4d01382..098fca0 100644 --- a/vibe/acp/tools/base.py +++ b/vibe/acp/tools/base.py @@ -4,12 +4,8 @@ from abc import abstractmethod from typing import Annotated, cast from acp import Client -from acp.helpers import ToolCallContentVariant -from acp.schema import ToolCallProgress from pydantic import BaseModel, ConfigDict, Field, SkipValidation -from vibe.acp.tools.session_update import resolve_kind -from vibe.core.logger import logger from vibe.core.tools.base import BaseTool, ToolError from vibe.core.tools.manager import ToolManager @@ -21,9 +17,6 @@ class AcpToolState(BaseModel): default=None, description="ACP Client" ) session_id: str | None = Field(default=None, description="Current ACP session ID") - tool_call_id: str | None = Field( - default=None, description="Current ACP tool call ID" - ) class BaseAcpTool[ToolState: AcpToolState](BaseTool): @@ -37,23 +30,17 @@ class BaseAcpTool[ToolState: AcpToolState](BaseTool): @classmethod def update_tool_state( - cls, - *, - tool_manager: ToolManager, - client: Client | None, - session_id: str | None, - tool_call_id: str | None, + cls, *, tool_manager: ToolManager, client: Client | None, session_id: str | None ) -> None: tool_instance = cls.get_tool_instance(cls.get_name(), tool_manager) tool_instance.state.client = client tool_instance.state.session_id = session_id - tool_instance.state.tool_call_id = tool_call_id @classmethod @abstractmethod def _get_tool_state_class(cls) -> type[ToolState]: ... - def _load_state(self) -> tuple[Client, str, str | None]: + def _load_state(self) -> tuple[Client, str]: if self.state.client is None: raise ToolError( "Client not available in tool state. This tool can only be used within an ACP session." @@ -63,27 +50,4 @@ class BaseAcpTool[ToolState: AcpToolState](BaseTool): "Session ID not available in tool state. This tool can only be used within an ACP session." ) - return self.state.client, self.state.session_id, self.state.tool_call_id - - async def _send_in_progress_session_update( - self, content: list[ToolCallContentVariant] | None = None - ) -> None: - client, session_id, tool_call_id = self._load_state() - if tool_call_id is None: - return - - tool_name = self.get_name() - try: - await client.session_update( - session_id=session_id, - update=ToolCallProgress( - session_update="tool_call_update", - tool_call_id=tool_call_id, - status="in_progress", - kind=resolve_kind(tool_name), - content=content, - field_meta={"tool_name": tool_name}, - ), - ) - except Exception as e: - logger.error(f"Failed to update session: {e!r}") + return self.state.client, self.state.session_id diff --git a/vibe/acp/tools/builtins/bash.py b/vibe/acp/tools/builtins/bash.py index bcdee5b..bdaaf15 100644 --- a/vibe/acp/tools/builtins/bash.py +++ b/vibe/acp/tools/builtins/bash.py @@ -6,7 +6,6 @@ from pathlib import Path from acp.schema import ( ContentToolCallContent, - TerminalToolCallContent, TextContentBlock, ToolCallProgress, ToolCallStart, @@ -15,6 +14,7 @@ from acp.schema import ( from vibe import VIBE_ROOT from vibe.acp.tools.base import AcpToolState, BaseAcpTool +from vibe.acp.tools.events import ToolTerminalOpenedEvent from vibe.acp.tools.session_update import resolve_kind from vibe.core.logger import logger from vibe.core.tools.base import BaseToolState, InvokeContext, ToolError @@ -37,7 +37,7 @@ class Bash(CoreBashTool, BaseAcpTool[AcpBashState]): async def run( self, args: BashArgs, ctx: InvokeContext | None = None ) -> AsyncGenerator[ToolStreamEvent | BashResult, None]: - client, session_id, _ = self._load_state() + client, session_id = self._load_state() timeout = args.timeout or self.config.default_timeout max_bytes = self.config.max_output_bytes @@ -54,11 +54,14 @@ class Bash(CoreBashTool, BaseAcpTool[AcpBashState]): terminal_id = terminal.terminal_id - await self._send_in_progress_session_update([ - TerminalToolCallContent(type="terminal", terminal_id=terminal_id) - ]) - try: + if ctx is not None: + yield ToolTerminalOpenedEvent( + tool_name=self.get_name(), + tool_call_id=ctx.tool_call_id, + terminal_id=terminal_id, + ) + exit_response = await self._wait_for_terminal_exit( terminal_id=terminal_id, timeout=timeout, command=args.command ) @@ -93,7 +96,7 @@ class Bash(CoreBashTool, BaseAcpTool[AcpBashState]): async def _wait_for_terminal_exit( self, terminal_id: str, timeout: int, command: str ) -> WaitForTerminalExitResponse: - client, session_id, _ = self._load_state() + client, session_id = self._load_state() try: return await asyncio.wait_for( diff --git a/vibe/acp/tools/builtins/read_file.py b/vibe/acp/tools/builtins/read_file.py index 69b0b45..b81785a 100644 --- a/vibe/acp/tools/builtins/read_file.py +++ b/vibe/acp/tools/builtins/read_file.py @@ -114,13 +114,11 @@ class ReadFile( ) async def _read_file(self, args: ReadFileArgs, file_path: Path) -> _ReadResult: - client, session_id, _ = self._load_state() + client, session_id = self._load_state() line = args.offset + 1 if args.offset > 0 else None limit = args.limit - await self._send_in_progress_session_update() - try: response = await client.read_text_file( session_id=session_id, path=str(file_path), line=line, limit=limit diff --git a/vibe/acp/tools/builtins/search_replace.py b/vibe/acp/tools/builtins/search_replace.py index a4c3d88..ff82698 100644 --- a/vibe/acp/tools/builtins/search_replace.py +++ b/vibe/acp/tools/builtins/search_replace.py @@ -44,9 +44,7 @@ class SearchReplace(CoreSearchReplaceTool, BaseAcpTool[AcpSearchReplaceState]): return AcpSearchReplaceState async def _read_file(self, file_path: Path) -> ReadSafeResult: - client, session_id, _ = self._load_state() - - await self._send_in_progress_session_update() + client, session_id = self._load_state() try: response = await client.read_text_file( @@ -74,7 +72,7 @@ class SearchReplace(CoreSearchReplaceTool, BaseAcpTool[AcpSearchReplaceState]): async def _write_file( self, file_path: Path, content: str, encoding: str, newline: str ) -> None: - client, session_id, _ = self._load_state() + client, session_id = self._load_state() try: await client.write_text_file( diff --git a/vibe/acp/tools/builtins/write_file.py b/vibe/acp/tools/builtins/write_file.py index f254544..3584d38 100644 --- a/vibe/acp/tools/builtins/write_file.py +++ b/vibe/acp/tools/builtins/write_file.py @@ -41,9 +41,7 @@ class WriteFile(CoreWriteFileTool, BaseAcpTool[AcpWriteFileState]): return AcpWriteFileState async def _write_file(self, args: WriteFileArgs, file_path: Path) -> None: - client, session_id, _ = self._load_state() - - await self._send_in_progress_session_update() + client, session_id = self._load_state() try: await client.write_text_file( diff --git a/vibe/acp/tools/events.py b/vibe/acp/tools/events.py new file mode 100644 index 0000000..0684212 --- /dev/null +++ b/vibe/acp/tools/events.py @@ -0,0 +1,8 @@ +from __future__ import annotations + +from vibe.core.types import ToolStreamEvent + + +class ToolTerminalOpenedEvent(ToolStreamEvent): + message: str = "" + terminal_id: str diff --git a/vibe/cli/textual_ui/app.py b/vibe/cli/textual_ui/app.py index a5121f6..79e208c 100644 --- a/vibe/cli/textual_ui/app.py +++ b/vibe/cli/textual_ui/app.py @@ -123,7 +123,10 @@ from vibe.core.agent_loop import AgentLoop, TeleportError from vibe.core.agents import AgentProfile from vibe.core.audio_player.audio_player import AudioPlayer from vibe.core.audio_recorder import AudioRecorder -from vibe.core.autocompletion.path_prompt import build_path_prompt_payload +from vibe.core.autocompletion.path_prompt import ( + build_path_prompt_payload, + build_title_segments, +) from vibe.core.autocompletion.path_prompt_adapter import render_path_prompt from vibe.core.config import VibeConfig from vibe.core.data_retention import DATA_RETENTION_MESSAGE @@ -140,6 +143,7 @@ from vibe.core.session.resume_sessions import ( ) from vibe.core.session.saved_sessions import update_saved_session_title_at_path from vibe.core.session.session_loader import SessionLoader +from vibe.core.session.title_format import format_session_title from vibe.core.skills.manager import SkillManager from vibe.core.teleport.telemetry import send_teleport_early_failure_telemetry from vibe.core.teleport.types import ( @@ -957,7 +961,7 @@ class VibeApp(App): # noqa: PLR0904 self.agent_loop.telemetry_client.send_slash_command_used(skill.name, "skill") prompt = SkillManager.build_skill_prompt(user_input, skill) - await self._handle_user_message(prompt) + await self._handle_user_message(prompt, title_source=user_input) return True @staticmethod @@ -1146,7 +1150,9 @@ class VibeApp(App): # noqa: PLR0904 return "\n\n".join(sections) - async def _handle_user_message(self, message: str) -> None: + async def _handle_user_message( + self, message: str, *, title_source: str | None = None + ) -> None: if self._remote_manager.is_active: await self._handle_remote_user_message(message) return @@ -1165,7 +1171,7 @@ class VibeApp(App): # noqa: PLR0904 await self._remote_manager.stop_stream() await self._remove_loading_widget() self._agent_task = asyncio.create_task( - self._handle_agent_loop_turn(message) + self._handle_agent_loop_turn(message, title_source=title_source) ) async def _handle_remote_user_message(self, message: str) -> None: @@ -1366,7 +1372,9 @@ class VibeApp(App): # noqa: PLR0904 event, loading_widget=self._loading_widget ) - async def _handle_agent_loop_turn(self, prompt: str) -> None: + async def _handle_agent_loop_turn( + self, prompt: str, *, title_source: str | None = None + ) -> None: self._agent_running = True await self._remove_loading_widget() @@ -1393,10 +1401,22 @@ class VibeApp(App): # noqa: PLR0904 message_id=message_id, ) rendered_prompt = render_path_prompt(prompt, base_dir=Path.cwd()) + auto_title: str | None = None + if self.agent_loop.session_logger.needs_initial_auto_title(): + auto_title = ( + format_session_title( + build_title_segments( + title_source or prompt, base_dir=Path.cwd() + ) + ) + or None + ) self._narrator_manager.cancel() self._narrator_manager.on_turn_start(rendered_prompt) async with aclosing( - self.agent_loop.act(rendered_prompt, client_message_id=message_id) + self.agent_loop.act( + rendered_prompt, client_message_id=message_id, auto_title=auto_title + ) ) as events: await self._handle_agent_loop_events(events) except asyncio.CancelledError: diff --git a/vibe/cli/textual_ui/handlers/event_handler.py b/vibe/cli/textual_ui/handlers/event_handler.py index 8716d56..604538c 100644 --- a/vibe/cli/textual_ui/handlers/event_handler.py +++ b/vibe/cli/textual_ui/handlers/event_handler.py @@ -33,6 +33,7 @@ from vibe.core.types import ( PlanReviewEndedEvent, PlanReviewRequestedEvent, ReasoningEvent, + SessionTitleUpdatedEvent, ToolCallEvent, ToolResultEvent, ToolStreamEvent, @@ -116,6 +117,8 @@ class EventHandler: case AgentProfileChangedEvent(): if self.on_profile_changed: self.on_profile_changed() + case SessionTitleUpdatedEvent(): + pass case UserMessageEvent(): await self.finalize_streaming() if self.is_remote: diff --git a/vibe/cli/textual_ui/widgets/mcp_app.py b/vibe/cli/textual_ui/widgets/mcp_app.py index 84694f1..1eb4315 100644 --- a/vibe/cli/textual_ui/widgets/mcp_app.py +++ b/vibe/cli/textual_ui/widgets/mcp_app.py @@ -71,6 +71,7 @@ _LIST_VIEW_HELP_AUTH = ( _DETAIL_VIEW_HELP = ( "↑↓ Navigate D Disable E Enable Backspace Back R Refresh Esc Close" ) +_DETAIL_VIEW_HELP_NO_TOOLS = "↑↓ Navigate Backspace Back R Refresh Esc Close" class MCPApp(Container): @@ -211,7 +212,16 @@ class MCPApp(Container): return self._status_message = "Refreshing..." - help = _DETAIL_VIEW_HELP if self._viewing_server else _LIST_VIEW_HELP_TOOLS + if self._viewing_server: + tools_source = ( + self._index.connector_tools + if self._viewing_kind == MCPSourceKind.CONNECTOR + else self._index.server_tools + ) + all_tools = tools_source.get(self._viewing_server, []) + help = _DETAIL_VIEW_HELP if all_tools else _DETAIL_VIEW_HELP_NO_TOOLS + else: + help = _LIST_VIEW_HELP_TOOLS self._set_help_text(help) self._refreshing = True @@ -516,9 +526,11 @@ class MCPApp(Container): self.query_one("#mcp-title", NoMarkupStatic).update( f"{title_prefix}: {server_name}" ) - self._set_help_text(_DETAIL_VIEW_HELP) tools_source = index.connector_tools if is_connector else index.server_tools all_tools = sorted(tools_source.get(server_name, []), key=lambda t: t[0]) + self._set_help_text( + _DETAIL_VIEW_HELP if all_tools else _DETAIL_VIEW_HELP_NO_TOOLS + ) if not all_tools: if ( is_connector diff --git a/vibe/core/agent_loop.py b/vibe/core/agent_loop.py index 7f092de..1dd9604 100644 --- a/vibe/core/agent_loop.py +++ b/vibe/core/agent_loop.py @@ -113,6 +113,7 @@ from vibe.core.types import ( RateLimitError, ReasoningEvent, Role, + SessionTitleUpdatedEvent, ToolCall, ToolCallEvent, ToolResultEvent, @@ -629,7 +630,11 @@ class AgentLoop: # noqa: PLR0904 @requires_init async def act( - self, msg: str, client_message_id: str | None = None + self, + msg: str, + client_message_id: str | None = None, + *, + auto_title: str | None = None, ) -> AsyncGenerator[BaseEvent, None]: self._clean_message_history() self.rewind_manager.create_checkpoint() @@ -639,7 +644,7 @@ class AgentLoop: # noqa: PLR0904 model_name = None async with agent_span(model=model_name, session_id=self.session_id): async for event in self._conversation_loop( - msg, client_message_id=client_message_id + msg, client_message_id=client_message_id, auto_title=auto_title ): yield event @@ -848,7 +853,11 @@ class AgentLoop: # noqa: PLR0904 return headers async def _conversation_loop( # noqa: PLR0912 - self, user_msg: str, client_message_id: str | None = None + self, + user_msg: str, + client_message_id: str | None = None, + *, + auto_title: str | None = None, ) -> AsyncGenerator[BaseEvent]: user_message = LLMMessage( role=Role.user, content=user_msg, message_id=client_message_id @@ -862,6 +871,11 @@ class AgentLoop: # noqa: PLR0904 yield UserMessageEvent(content=user_msg, message_id=user_message.message_id) + if auto_title is not None and self.session_logger.set_initial_auto_title( + auto_title + ): + yield SessionTitleUpdatedEvent(title=auto_title) + if self._hooks_manager: self._hooks_manager.reset_retry_count() diff --git a/vibe/core/autocompletion/path_prompt.py b/vibe/core/autocompletion/path_prompt.py index f0b5b2e..996b0b9 100644 --- a/vibe/core/autocompletion/path_prompt.py +++ b/vibe/core/autocompletion/path_prompt.py @@ -4,6 +4,8 @@ from dataclasses import dataclass from pathlib import Path from typing import Literal +from vibe.core.session.title_format import MentionSegment, TextSegment, TitleSegment + @dataclass(frozen=True, slots=True) class PathResource: @@ -107,3 +109,35 @@ def _dedupe_resources(resources: list[PathResource]) -> list[PathResource]: seen.add(resource.path) unique.append(resource) return unique + + +def build_title_segments( + message: str, *, base_dir: Path | None = None +) -> list[TitleSegment]: + if not message: + return [] + + resolved_base = (base_dir or Path.cwd()).resolve() + segments: list[TitleSegment] = [] + text_buf: list[str] = [] + pos = 0 + + def flush_text() -> None: + if text_buf: + segments.append(TextSegment(text="".join(text_buf))) + text_buf.clear() + + while pos < len(message): + if _is_path_anchor(message, pos): + candidate, new_pos = _extract_candidate(message, pos + 1) + if candidate and (resource := _to_resource(candidate, resolved_base)): + flush_text() + segments.append(MentionSegment(name=resource.path.name)) + pos = new_pos + continue + + text_buf.append(message[pos]) + pos += 1 + + flush_text() + return segments diff --git a/vibe/core/config/builder.py b/vibe/core/config/builder.py new file mode 100644 index 0000000..a75fc35 --- /dev/null +++ b/vibe/core/config/builder.py @@ -0,0 +1,122 @@ +from __future__ import annotations + +import asyncio +from collections import defaultdict +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any + +from vibe.core.config.layer import ( + ConfigLayer, + EmptyLayerError, + RawConfig, + UntrustedLayerError, +) +from vibe.core.config.schema import ConfigFragment, ConfigSchema, MergeFieldMetadata + + +@dataclass(frozen=True, slots=True) +class _LayerData: + name: str + data: dict[str, Any] + + +class ConfigBuilder[S: ConfigSchema]: + """Collects layers and merges them into an immutable Config[S].""" + + def __init__(self, schema: type[S]) -> None: + self._schema = schema + self._layers: list[ConfigLayer[RawConfig]] = [] + self._lock = asyncio.Lock() + + def add_layer(self, layer: ConfigLayer[RawConfig]) -> None: + self._layers.append(layer) + + def add_layers(self, layers: list[ConfigLayer[RawConfig]]) -> None: + self._layers.extend(layers) + + @property + def layers(self) -> list[ConfigLayer[RawConfig]]: + return self._layers + + async def build(self, force_load: bool = False) -> S: + """Merge all layers and return a validated schema. + + Untrusted and empty layers are skipped. + Pass ``force_load=True`` to bypass caching. + """ + async with self._lock: + internal_layers = self._layers.copy() + + layer_dicts: list[_LayerData] = [] + for layer in internal_layers: + try: + data = await layer.load(force=force_load) + raw = data.model_dump() + if raw: + layer_dicts.append(_LayerData(name=layer.name, data=raw)) + except (UntrustedLayerError, EmptyLayerError): + continue + + merged, origins = self._merge_fields(self._schema, layer_dicts) + return self._schema(origins=origins, **merged) + + def _merge_fields( + self, schema: type[S], layer_dicts: list[_LayerData] + ) -> tuple[dict[str, Any], dict[str, Any]]: + accumulated: dict[str, Any] = defaultdict(dict) + origins: dict[str, Any] = {} + + for ld in layer_dicts: + for key, value in ld.data.items(): + if key not in schema.model_fields: + continue + + field_info = schema.model_fields[key] + annotation = field_info.annotation + if annotation is None: + continue + + is_fragment = isinstance(annotation, type) and issubclass( + annotation, ConfigFragment + ) + if is_fragment: + if not isinstance(value, dict): + continue + + for fragment_key, fragment_value in value.items(): + if fragment_key not in annotation.model_fields: + continue + + fragment_field = annotation.model_fields[fragment_key] + fragment_meta = MergeFieldMetadata.from_field(fragment_field) + if fragment_meta is None: + continue + + accumulated[key][fragment_key] = ( + fragment_meta.merge_strategy.apply( + accumulated[key].get(fragment_key), + fragment_value, + key_fn=self._make_key_fn(fragment_meta), + ) + ) + continue + + meta = MergeFieldMetadata.from_field(field_info) + if meta is None: + continue + + accumulated[key] = meta.merge_strategy.apply( + accumulated.get(key), value, key_fn=self._make_key_fn(meta) + ) + + return accumulated, origins + + def _make_key_fn( + self, merge_field_meta: MergeFieldMetadata + ) -> Callable[[Any], str] | None: + merge_key = merge_field_meta.merge_key + if merge_key is None: + return None + + return lambda item: item[merge_key] diff --git a/vibe/core/config/layers/__init__.py b/vibe/core/config/layers/__init__.py new file mode 100644 index 0000000..e79996f --- /dev/null +++ b/vibe/core/config/layers/__init__.py @@ -0,0 +1,5 @@ +from __future__ import annotations + +from vibe.core.config.layers.user import UserConfigLayer + +__all__ = ["UserConfigLayer"] diff --git a/vibe/core/config/layers/user.py b/vibe/core/config/layers/user.py new file mode 100644 index 0000000..d88a197 --- /dev/null +++ b/vibe/core/config/layers/user.py @@ -0,0 +1,32 @@ +from __future__ import annotations + +from pathlib import Path +import tomllib +from typing import Any + +from vibe.core.config.layer import ConfigLayer, RawConfig +from vibe.core.paths._vibe_home import VIBE_HOME + + +class UserConfigLayer(ConfigLayer[RawConfig]): + """Reads the user-level TOML config file. Always trusted. + + Defaults to ``~/.vibe/config.toml`` (via VIBE_HOME). + Pass an explicit ``path`` for testing. + """ + + def __init__(self, *, path: Path | None = None, name: str = "user-toml") -> None: + super().__init__(name=name) + self._path = path or (VIBE_HOME.path / "config.toml") + + async def _check_trust(self) -> bool: + return True + + async def _read_config(self) -> dict[str, Any]: + if not self._path.exists(): + return {} + with self._path.open("rb") as f: + return tomllib.load(f) + + async def apply(self, patch: Any, *, on_conflict: str = "cancel") -> None: + raise NotImplementedError("UserConfigLayer.apply() is not implemented (M2)") diff --git a/vibe/core/config/orchestrator.py b/vibe/core/config/orchestrator.py new file mode 100644 index 0000000..b4af9c4 --- /dev/null +++ b/vibe/core/config/orchestrator.py @@ -0,0 +1,46 @@ +from __future__ import annotations + +from typing import Any + +from vibe.core.config.builder import ConfigBuilder +from vibe.core.config.layer import ConfigLayer, RawConfig +from vibe.core.config.schema import ConfigSchema + + +class ConfigOrchestrator[S: ConfigSchema]: + """Single entry point for config management.""" + + def __init__(self, builder: ConfigBuilder[S], config: S) -> None: + self._builder = builder + self._config = config + + @classmethod + async def create( + cls, *, schema: type[S], layers: list[ConfigLayer[RawConfig]] + ) -> ConfigOrchestrator[S]: + """Build an orchestrator from a schema and an ordered list of layers.""" + builder = ConfigBuilder[S](schema) + builder.add_layers(layers) + config = await builder.build() + instance = cls(builder, config) + return instance + + @property + def config(self) -> S: + return self._config + + def get_layer(self, name: str) -> ConfigLayer[RawConfig]: + for layer in self._builder.layers: + if layer.name == name: + return layer + raise KeyError(f"No layer named {name!r}") + + async def reload(self) -> None: + """Force-reload all layers and atomically replace the config snapshot.""" + self._config = await self._builder.build(force_load=True) + + async def apply_patch(self, patch: Any) -> None: + raise NotImplementedError("apply_patch() is not implemented (M2)") + + async def subscribe(self, callback: Any) -> None: + raise NotImplementedError("subscribe() is not implemented (M3)") diff --git a/vibe/core/config/schema.py b/vibe/core/config/schema.py index 10479d0..8c2217e 100644 --- a/vibe/core/config/schema.py +++ b/vibe/core/config/schema.py @@ -2,7 +2,7 @@ from __future__ import annotations from dataclasses import dataclass, field -from pydantic import BaseModel +from pydantic import BaseModel, ConfigDict, PrivateAttr from pydantic.fields import FieldInfo from vibe.core.utils.merge import MergeStrategy @@ -19,11 +19,25 @@ class ConfigDefinitionError(TypeError): class ConfigSchema(BaseModel): """Base for composite config schemas composed of fragments and merge-aware fields.""" + model_config = ConfigDict(frozen=True) + + _origins: dict[str, str] = PrivateAttr(default_factory=dict) + + def __init__( + self, *, origins: dict[str, str] | None = None, **data: object + ) -> None: + super().__init__(**data) + if origins: + self._origins = origins + + def origin_of(self, key: str) -> str | None: + return self._origins.get(key) + @classmethod def __pydantic_on_complete__(cls) -> None: super().__pydantic_on_complete__() - if cls.__name__ == "ConfigSchema" and cls.__module__ == __name__: + if not cls.model_fields: return for field_name, field_info in cls.model_fields.items(): @@ -50,6 +64,8 @@ class ConfigSchema(BaseModel): class ConfigFragment(BaseModel): """Base for domain config groups with merge-aware top-level fields.""" + model_config = ConfigDict(frozen=True) + @classmethod def __pydantic_on_complete__(cls) -> None: super().__pydantic_on_complete__() diff --git a/vibe/core/config/types.py b/vibe/core/config/types.py new file mode 100644 index 0000000..5644beb --- /dev/null +++ b/vibe/core/config/types.py @@ -0,0 +1,8 @@ +from __future__ import annotations + +from dataclasses import dataclass + + +@dataclass(frozen=True, slots=True) +class ConfigOrigin: + layer_name: str diff --git a/vibe/core/session/session_logger.py b/vibe/core/session/session_logger.py index fbcb858..7ff73a7 100644 --- a/vibe/core/session/session_logger.py +++ b/vibe/core/session/session_logger.py @@ -18,6 +18,7 @@ from vibe.core.session.session_loader import ( METADATA_FILENAME, SessionLoader, ) +from vibe.core.session.title_format import MAX_TITLE_LENGTH from vibe.core.types import AgentStats, LLMMessage, Role, SessionMetadata from vibe.core.utils import is_windows, utc_now from vibe.core.utils.io import read_safe_async @@ -145,7 +146,7 @@ class SessionLogger: title_source="auto", ) - def _get_title(self, messages: Sequence[LLMMessage]) -> str: + def _fallback_title_from_messages(self, messages: Sequence[LLMMessage]) -> str: first_user_message = None for message in messages: if message.role == Role.user: @@ -153,14 +154,12 @@ class SessionLogger: break if first_user_message is None: - title = "Untitled session" - else: - MAX_TITLE_LENGTH = 50 - text = str(first_user_message.content) - title = text[:MAX_TITLE_LENGTH] - if len(text) > MAX_TITLE_LENGTH: - title += "…" + return "Untitled session" + text = str(first_user_message.content) + title = text[:MAX_TITLE_LENGTH] + if len(text) > MAX_TITLE_LENGTH: + title += "…" return title def _set_title_state( @@ -183,14 +182,28 @@ class SessionLogger: self._set_title_state(normalized_title, source="manual") + def needs_initial_auto_title(self) -> bool: + return self.session_metadata is not None and self.session_metadata.title is None + + def set_initial_auto_title(self, title: str) -> bool: + if not self.needs_initial_auto_title(): + return False + + normalized_title = title.strip() + if not normalized_title: + return False + + self._set_title_state(normalized_title, source="auto") + return True + def _resolve_title(self, messages: Sequence[LLMMessage]) -> str | None: if self.session_metadata is None: - return self._get_title(messages) + return self._fallback_title_from_messages(messages) - if self.session_metadata.title_source == "manual": + if self.session_metadata.title is not None: return self.session_metadata.title - title = self._get_title(messages) + title = self._fallback_title_from_messages(messages) self._set_title_state(title, source="auto") return title diff --git a/vibe/core/session/title_format.py b/vibe/core/session/title_format.py new file mode 100644 index 0000000..10ceb5c --- /dev/null +++ b/vibe/core/session/title_format.py @@ -0,0 +1,51 @@ +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass +import re + +MAX_TITLE_LENGTH = 50 +ELLIPSIS = "…" + +_WHITESPACE_RE = re.compile(r"\s+") + + +@dataclass(frozen=True, slots=True) +class TextSegment: + text: str + + +@dataclass(frozen=True, slots=True) +class MentionSegment: + name: str + start_line: int | None = None + end_line: int | None = None + + +TitleSegment = TextSegment | MentionSegment + + +def format_session_title(segments: Sequence[TitleSegment]) -> str: + parts = [_render_segment(s) for s in segments] + joined = "".join(parts) + collapsed = _WHITESPACE_RE.sub(" ", joined).strip() + + if not collapsed: + return "" + + if len(collapsed) > MAX_TITLE_LENGTH: + return collapsed[:MAX_TITLE_LENGTH] + ELLIPSIS + + return collapsed + + +def _render_segment(segment: TitleSegment) -> str: + match segment: + case TextSegment(text=text): + return text + case MentionSegment(name=name, start_line=start, end_line=end): + if start is not None and end is not None: + return f"@{name}:{start}-{end}" + if start is not None: + return f"@{name}:{start}" + return f"@{name}" diff --git a/vibe/core/types.py b/vibe/core/types.py index cf87ec7..d63f11a 100644 --- a/vibe/core/types.py +++ b/vibe/core/types.py @@ -456,6 +456,10 @@ class AgentProfileChangedEvent(BaseEvent): agent_name: str +class SessionTitleUpdatedEvent(BaseEvent): + title: str + + class OutputFormat(StrEnum): TEXT = auto() JSON = auto() diff --git a/vibe/setup/auth/__init__.py b/vibe/setup/auth/__init__.py index f9d5003..b8f7e18 100644 --- a/vibe/setup/auth/__init__.py +++ b/vibe/setup/auth/__init__.py @@ -1,6 +1,10 @@ from __future__ import annotations -from vibe.setup.auth.browser_sign_in import BrowserSignInService, BrowserSignInStatus +from vibe.setup.auth.browser_sign_in import ( + BrowserSignInAttempt, + BrowserSignInService, + BrowserSignInStatus, +) from vibe.setup.auth.browser_sign_in_gateway import ( BrowserSignInError, BrowserSignInErrorCode, @@ -11,6 +15,7 @@ from vibe.setup.auth.browser_sign_in_gateway import ( from vibe.setup.auth.http_browser_sign_in_gateway import HttpBrowserSignInGateway __all__ = [ + "BrowserSignInAttempt", "BrowserSignInError", "BrowserSignInErrorCode", "BrowserSignInGateway", diff --git a/vibe/setup/auth/api_key_persistence.py b/vibe/setup/auth/api_key_persistence.py new file mode 100644 index 0000000..681b7be --- /dev/null +++ b/vibe/setup/auth/api_key_persistence.py @@ -0,0 +1,64 @@ +from __future__ import annotations + +import os + +from dotenv import set_key + +from vibe.core.config import DEFAULT_PROVIDERS, ProviderConfig, VibeConfig +from vibe.core.paths import GLOBAL_ENV_FILE +from vibe.core.telemetry.send import TelemetryClient +from vibe.core.telemetry.types import EntrypointMetadata +from vibe.core.types import Backend + + +def _save_api_key_to_env_file(env_key: str, api_key: str) -> None: + GLOBAL_ENV_FILE.path.parent.mkdir(parents=True, exist_ok=True) + set_key(GLOBAL_ENV_FILE.path, env_key, api_key) + + +def _get_mistral_provider() -> ProviderConfig: + return next( + provider for provider in DEFAULT_PROVIDERS if provider.name == "mistral" + ) + + +def _load_onboarding_provider() -> ProviderConfig: + from vibe.setup.onboarding.context import OnboardingContext + + return OnboardingContext.load().provider + + +def resolve_api_key_provider(provider: ProviderConfig | None = None) -> ProviderConfig: + resolved_provider = provider or _load_onboarding_provider() + if resolved_provider.api_key_env_var: + return resolved_provider + return _get_mistral_provider() + + +def persist_api_key( + provider: ProviderConfig, + api_key: str, + *, + entrypoint_metadata: EntrypointMetadata | None = None, +) -> str: + env_key = provider.api_key_env_var + if not env_key: + return "env_var_error:" + try: + os.environ[env_key] = api_key + except ValueError: + return f"env_var_error:{env_key}" + try: + _save_api_key_to_env_file(env_key, api_key) + except (OSError, ValueError) as err: + return f"save_error:{err}" + if provider.backend == Backend.MISTRAL: + try: + telemetry = TelemetryClient( + config_getter=VibeConfig, + entrypoint_metadata_getter=lambda: entrypoint_metadata, + ) + telemetry.send_onboarding_api_key_added() + except Exception: + pass + return "completed" diff --git a/vibe/setup/auth/browser_sign_in.py b/vibe/setup/auth/browser_sign_in.py index 3cb2672..e024c29 100644 --- a/vibe/setup/auth/browser_sign_in.py +++ b/vibe/setup/auth/browser_sign_in.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio import base64 from collections.abc import Awaitable, Callable +from dataclasses import dataclass from datetime import UTC, datetime from enum import StrEnum import hashlib @@ -13,7 +14,6 @@ from vibe.setup.auth.browser_sign_in_gateway import ( BrowserSignInError, BrowserSignInErrorCode, BrowserSignInGateway, - BrowserSignInProcess, ) @@ -30,6 +30,15 @@ SleepFn = Callable[[float], Awaitable[None]] NowFn = Callable[[], datetime] +@dataclass(frozen=True) +class BrowserSignInAttempt: + process_id: str + sign_in_url: str + poll_url: str + expires_at: datetime + code_verifier: str + + class BrowserSignInService: _max_consecutive_poll_failures = 3 @@ -51,38 +60,55 @@ class BrowserSignInService: async def aclose(self) -> None: await self._gateway.aclose() - async def authenticate(self, status_callback: StatusCallback | None = None) -> str: + async def start_attempt(self) -> BrowserSignInAttempt: verifier, challenge = _generate_pkce_pair() process = await self._gateway.create_process(challenge) - self._emit(status_callback, BrowserSignInStatus.OPENING_BROWSER) - self._open_browser_or_raise(process.sign_in_url) + return BrowserSignInAttempt( + process_id=process.process_id, + sign_in_url=process.sign_in_url, + poll_url=process.poll_url, + expires_at=process.expires_at, + code_verifier=verifier, + ) + + async def complete_attempt( + self, + attempt: BrowserSignInAttempt, + status_callback: StatusCallback | None = None, + ) -> str: self._emit(status_callback, BrowserSignInStatus.WAITING_FOR_BROWSER_SIGN_IN) - exchange_token = await self._wait_for_completion(process) + exchange_token = await self._wait_for_completion(attempt) self._emit(status_callback, BrowserSignInStatus.EXCHANGING) api_key = await self._gateway.exchange( - process.process_id, exchange_token, verifier + attempt.process_id, exchange_token, attempt.code_verifier ) self._emit(status_callback, BrowserSignInStatus.COMPLETED) return api_key - async def _wait_for_completion(self, process: BrowserSignInProcess) -> str: + async def authenticate(self, status_callback: StatusCallback | None = None) -> str: + attempt = await self.start_attempt() + self._emit(status_callback, BrowserSignInStatus.OPENING_BROWSER) + self._open_browser_or_raise(attempt.sign_in_url) + return await self.complete_attempt(attempt, status_callback=status_callback) + + async def _wait_for_completion(self, attempt: BrowserSignInAttempt) -> str: consecutive_poll_failures = 0 - while self._now() < process.expires_at: + while self._now() < attempt.expires_at: try: - payload = await self._gateway.poll(process.poll_url) + payload = await self._gateway.poll(attempt.poll_url) except BrowserSignInError as err: if err.code is not BrowserSignInErrorCode.POLL_FAILED: raise consecutive_poll_failures += 1 if consecutive_poll_failures >= self._max_consecutive_poll_failures: raise - await self._sleep_until_next_poll_or_timeout(process.expires_at) + await self._sleep_until_next_poll_or_timeout(attempt.expires_at) continue consecutive_poll_failures = 0 match payload.status: case "pending": - await self._sleep_until_next_poll_or_timeout(process.expires_at) + await self._sleep_until_next_poll_or_timeout(attempt.expires_at) case "completed": if payload.exchange_token: return payload.exchange_token diff --git a/vibe/setup/auth/http_browser_sign_in_gateway.py b/vibe/setup/auth/http_browser_sign_in_gateway.py index 5499ade..b92616f 100644 --- a/vibe/setup/auth/http_browser_sign_in_gateway.py +++ b/vibe/setup/auth/http_browser_sign_in_gateway.py @@ -188,9 +188,21 @@ class HttpBrowserSignInGateway(BrowserSignInGateway): json={"exchange_token": exchange_token, "code_verifier": code_verifier}, ) except httpx.HTTPError as err: + logger.warning( + "Browser sign-in exchange request failed for api_base_url=%s: %s", + self._api_base_url, + err, + exc_info=True, + ) raise BrowserSignInError(message, code=code) from err if not response.is_success: + logger.warning( + "Browser sign-in exchange request returned status_code=%s for api_base_url=%s response_detail=%s", + response.status_code, + self._api_base_url, + _build_safe_response_error_detail(response), + ) raise BrowserSignInError(message, code=code) raw_payload = _response_json_or_raise(response, message=message, code=code) @@ -298,6 +310,29 @@ def _response_json_or_raise( return dict(payload) +def _build_safe_response_error_detail(response: httpx.Response) -> str: + try: + payload = response.json() + except ValueError: + return "unavailable" + + if not isinstance(payload, Mapping): + return "unavailable" + + for key in ("detail", "message", "error"): + value = payload.get(key) + if isinstance(value, str): + return _truncate_log_value(value) + + return "unavailable" + + +def _truncate_log_value(value: str, *, max_length: int = 200) -> str: + if len(value) <= max_length: + return value + return f"{value[: max_length - 3]}..." + + def _build_safe_url_log_details(value: object) -> str: if not isinstance(value, str): return "unavailable" diff --git a/vibe/setup/onboarding/screens/api_key.py b/vibe/setup/onboarding/screens/api_key.py index b02a405..40b87c6 100644 --- a/vibe/setup/onboarding/screens/api_key.py +++ b/vibe/setup/onboarding/screens/api_key.py @@ -1,9 +1,7 @@ from __future__ import annotations -import os from typing import ClassVar -from dotenv import set_key from textual.app import ComposeResult from textual.binding import Binding, BindingType from textual.containers import Center, Horizontal, Vertical @@ -13,13 +11,13 @@ from textual.widgets import Input, Link, Static from vibe.cli.clipboard import copy_selection_to_clipboard from vibe.cli.textual_ui.widgets.no_markup_static import NoMarkupStatic -from vibe.core.config import DEFAULT_PROVIDERS, ProviderConfig, VibeConfig -from vibe.core.paths import GLOBAL_ENV_FILE -from vibe.core.telemetry.send import TelemetryClient +from vibe.core.config import ProviderConfig from vibe.core.telemetry.types import EntrypointMetadata -from vibe.core.types import Backend +from vibe.setup.auth.api_key_persistence import ( + persist_api_key, + resolve_api_key_provider, +) from vibe.setup.onboarding.base import OnboardingScreen -from vibe.setup.onboarding.context import OnboardingContext PROVIDER_HELP = { "mistral": ("https://console.mistral.ai/codestral/cli", "Mistral AI Studio") @@ -29,55 +27,6 @@ CONFIG_DOCS_URL = ( ) -def _save_api_key_to_env_file(env_key: str, api_key: str) -> None: - GLOBAL_ENV_FILE.path.parent.mkdir(parents=True, exist_ok=True) - set_key(GLOBAL_ENV_FILE.path, env_key, api_key) - - -def persist_api_key( - provider: ProviderConfig, - api_key: str, - *, - entrypoint_metadata: EntrypointMetadata | None = None, -) -> str: - env_key = provider.api_key_env_var - if not env_key: - return "env_var_error:" - try: - os.environ[env_key] = api_key - except ValueError: - return f"env_var_error:{env_key}" - try: - _save_api_key_to_env_file(env_key, api_key) - except (OSError, ValueError) as err: - return f"save_error:{err}" - if provider.backend == Backend.MISTRAL: - try: - telemetry = TelemetryClient( - config_getter=VibeConfig, - entrypoint_metadata_getter=lambda: entrypoint_metadata, - ) - telemetry.send_onboarding_api_key_added() - except Exception: - pass - return "completed" - - -def _get_mistral_provider() -> ProviderConfig: - return next( - provider for provider in DEFAULT_PROVIDERS if provider.name == "mistral" - ) - - -def _resolve_onboarding_provider( - provider: ProviderConfig | None = None, -) -> ProviderConfig: - resolved_provider = provider or OnboardingContext.load().provider - if resolved_provider.api_key_env_var: - return resolved_provider - return _get_mistral_provider() - - class ApiKeyScreen(OnboardingScreen): BINDINGS: ClassVar[list[BindingType]] = [ Binding("ctrl+c", "cancel", "Cancel", show=False), @@ -93,7 +42,7 @@ class ApiKeyScreen(OnboardingScreen): entrypoint_metadata: EntrypointMetadata | None = None, ) -> None: super().__init__() - self.provider = _resolve_onboarding_provider(provider) + self.provider = resolve_api_key_provider(provider) self._entrypoint_metadata = entrypoint_metadata def _compose_provider_link(self, provider_name: str) -> ComposeResult: diff --git a/vibe/setup/onboarding/screens/browser_sign_in.py b/vibe/setup/onboarding/screens/browser_sign_in.py index 2eb967d..e1a360c 100644 --- a/vibe/setup/onboarding/screens/browser_sign_in.py +++ b/vibe/setup/onboarding/screens/browser_sign_in.py @@ -22,11 +22,11 @@ from vibe.setup.auth import ( BrowserSignInService, BrowserSignInStatus, ) -from vibe.setup.onboarding.base import OnboardingScreen -from vibe.setup.onboarding.screens.api_key import ( - _resolve_onboarding_provider, +from vibe.setup.auth.api_key_persistence import ( persist_api_key, + resolve_api_key_provider, ) +from vibe.setup.onboarding.base import OnboardingScreen PENDING_HINT = "Press M to enter API key manually · Esc to cancel" ERROR_HINT = "Press R to retry · Press M to enter API key manually · Esc to cancel" @@ -215,7 +215,7 @@ class BrowserSignInScreen(OnboardingScreen): self._worker = None self.app.exit( persist_api_key( - _resolve_onboarding_provider(self.provider), + resolve_api_key_provider(self.provider), api_key, entrypoint_metadata=self._entrypoint_metadata, )