Co-authored-by: Antoine <33425718+anth2o@users.noreply.github.com>
Co-authored-by: Bastien <bastien.baret@gmail.com>
Co-authored-by: Clément Sirieix <clement.sirieix@mistral.ai>
Co-authored-by: Kim-Adeline Miguel <51720070+kimadeline@users.noreply.github.com>
Co-authored-by: Mathias Gesbert <mathias.gesbert@mistral.ai>
Co-authored-by: Maxime Dolores <maxime.dolores@ext.mistral.ai>
Co-authored-by: Michel Thomazo <51709227+michelTho@users.noreply.github.com>
Co-authored-by: Nelson PROIA <144663685+Nelson-PROIA@users.noreply.github.com>
Co-authored-by: Pierre Rossinès <pierre.rossines@mistral.ai>
Co-authored-by: Quentin <quentin.torroba@mistral.ai>
Co-authored-by: Robin Gullo <robin.gullo@mistral.ai>
Co-authored-by: Mistral Vibe <vibe@mistral.ai>
This commit is contained in:
Clément Drouin 2026-04-28 17:44:07 +02:00 committed by GitHub
parent a83c81ecf5
commit 632ea8c032
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
253 changed files with 13965 additions and 2525 deletions

View file

@ -1,8 +1,10 @@
from __future__ import annotations
import asyncio
from collections.abc import AsyncGenerator
from collections.abc import AsyncGenerator, Callable
from contextlib import aclosing
import inspect
import logging
import os
from pathlib import Path
import sys
@ -45,8 +47,10 @@ from acp.schema import (
PromptCapabilities,
ResumeSessionResponse,
SessionCapabilities,
SessionCloseCapabilities,
SessionConfigOptionBoolean,
SessionConfigOptionSelect,
SessionForkCapabilities,
SessionInfo,
SessionListCapabilities,
SetSessionConfigOptionResponse,
@ -60,13 +64,14 @@ from acp.schema import (
Usage,
UsageUpdate,
)
from pydantic import BaseModel, ConfigDict, SkipValidation
from pydantic import AliasChoices, BaseModel, ConfigDict, Field, ValidationError
from vibe import VIBE_ROOT, __version__
from vibe.acp.acp_logger import acp_message_observer
from vibe.acp.commands import AcpCommandRegistry
from vibe.acp.exceptions import (
ConfigurationError,
ContextTooLongError,
ConversationLimitError,
InternalError,
InvalidRequestError,
@ -76,13 +81,18 @@ from vibe.acp.exceptions import (
SessionNotFoundError,
UnauthenticatedError,
)
from vibe.acp.session import AcpSessionLoop
from vibe.acp.tools.base import BaseAcpTool
from vibe.acp.tools.session_update import (
tool_call_session_update,
tool_result_session_update,
)
from vibe.acp.utils import (
THINKING_LEVELS,
ThinkingLevel,
ToolOption,
build_mode_state,
build_model_state,
build_permission_options,
create_assistant_message_replay,
create_compact_end_session_update,
@ -93,8 +103,7 @@ from vibe.acp.utils import (
create_user_message_replay,
get_proxy_help_text,
is_valid_acp_mode,
make_mode_response,
make_model_response,
make_thinking_response,
)
from vibe.core.agent_loop import AgentLoop
from vibe.core.agents.models import CHAT as CHAT_AGENT, BuiltinAgentName
@ -106,6 +115,7 @@ from vibe.core.config import (
load_dotenv_values,
)
from vibe.core.data_retention import DATA_RETENTION_MESSAGE
from vibe.core.hooks.config import load_hooks_from_fs
from vibe.core.proxy_setup import (
ProxySetupError,
parse_proxy_command,
@ -114,6 +124,9 @@ from vibe.core.proxy_setup import (
)
from vibe.core.session.session_loader import SessionLoader
from vibe.core.skills.manager import SkillManager
from vibe.core.telemetry.build_metadata import build_entrypoint_metadata
from vibe.core.telemetry.send import TelemetryClient
from vibe.core.telemetry.types import EntrypointMetadata
from vibe.core.tools.permissions import RequiredPermission
from vibe.core.types import (
AgentProfileChangedEvent,
@ -122,7 +135,7 @@ from vibe.core.types import (
AssistantEvent,
CompactEndEvent,
CompactStartEvent,
EntrypointMetadata,
ContextTooLongError as CoreContextTooLongError,
LLMMessage,
RateLimitError as CoreRateLimitError,
ReasoningEvent,
@ -137,6 +150,25 @@ from vibe.core.utils import (
get_user_cancellation_message,
)
logger = logging.getLogger("vibe")
class ForkSessionParams(BaseModel):
model_config = ConfigDict(extra="ignore", populate_by_name=True)
message_id: str | None = Field(default=None, alias="messageId")
class TelemetrySendNotification(BaseModel):
model_config = ConfigDict(extra="ignore")
event: str
properties: dict[str, Any] = Field(default_factory=dict)
session_id: str = Field(validation_alias=AliasChoices("session_id", "sessionId"))
_EVENT_DISPATCHERS: dict[str, Callable[[TelemetryClient, dict[str, Any]], None]] = {}
def _resolved_user_message_id(client_message_id: str | None) -> str:
if client_message_id is not None:
@ -144,14 +176,6 @@ def _resolved_user_message_id(client_message_id: str | None) -> str:
return str(uuid4())
class AcpSessionLoop(BaseModel):
model_config = ConfigDict(arbitrary_types_allowed=True)
id: str
agent_loop: AgentLoop
command_registry: SkipValidation[AcpCommandRegistry]
task: asyncio.Task[None] | None = None
class VibeAcpAgentLoop(AcpAgent):
client: Client
@ -220,7 +244,9 @@ class VibeAcpAgentLoop(AcpAgent):
audio=False, embedded_context=True, image=False
),
session_capabilities=SessionCapabilities(
list=SessionListCapabilities()
close=SessionCloseCapabilities(),
list=SessionListCapabilities(),
fork=SessionForkCapabilities(),
),
),
protocol_version=PROTOCOL_VERSION,
@ -240,7 +266,7 @@ class VibeAcpAgentLoop(AcpAgent):
raise NotImplementedMethodError("authenticate")
def _build_entrypoint_metadata(self) -> EntrypointMetadata:
return EntrypointMetadata(
return build_entrypoint_metadata(
agent_entrypoint="acp",
agent_version=__version__,
client_name=self.client_info.name if self.client_info else "",
@ -266,17 +292,44 @@ class VibeAcpAgentLoop(AcpAgent):
)
self.sessions[session.id] = session
command_registry.set_on_changed(
lambda: self._send_available_commands(session.id)
)
async def _on_commands_changed() -> None:
session.spawn(self._send_available_commands(session))
if not agent_loop.auto_approve:
command_registry.set_on_changed(_on_commands_changed)
if not agent_loop.bypass_tool_permissions:
agent_loop.set_approval_callback(self._create_approval_callback(session.id))
asyncio.create_task(self._send_available_commands(session.id))
session.spawn(self._send_available_commands(session))
return session
def _create_agent_loop(
self, config: VibeConfig, agent_name: str, hook_config_result: Any = None
) -> AgentLoop:
agent_loop = AgentLoop(
config=config,
agent_name=agent_name,
enable_streaming=True,
entrypoint_metadata=self._build_entrypoint_metadata(),
defer_heavy_init=True,
hook_config_result=hook_config_result,
)
agent_loop.agent_manager.register_agent(CHAT_AGENT)
return agent_loop
def _build_session_state(
self, session: AcpSessionLoop
) -> tuple[Any, Any, Any, Any]:
modes_state, modes_config = build_mode_state(
list(session.agent_loop.agent_manager.available_agents.values()),
session.agent_loop.agent_profile.name,
)
models_state, models_config = build_model_state(
session.agent_loop.config.models, session.agent_loop.config.active_model
)
return modes_state, modes_config, models_state, models_config
@override
async def new_session(
self,
@ -288,16 +341,12 @@ class VibeAcpAgentLoop(AcpAgent):
os.chdir(cwd)
config = self._load_config()
hook_config_result = load_hooks_from_fs(config)
try:
agent_loop = AgentLoop(
config=config,
agent_name=BuiltinAgentName.DEFAULT,
enable_streaming=True,
entrypoint_metadata=self._build_entrypoint_metadata(),
defer_heavy_init=True,
agent_loop = self._create_agent_loop(
config, BuiltinAgentName.DEFAULT, hook_config_result=hook_config_result
)
agent_loop.agent_manager.register_agent(CHAT_AGENT)
# NOTE: For now, we pin session.id to agent_loop.session_id right after init time.
# We should just use agent_loop.session_id everywhere, but it can still change during
# session lifetime (e.g. agent_loop.compact is called).
@ -308,19 +357,13 @@ class VibeAcpAgentLoop(AcpAgent):
agent_loop.emit_new_session_telemetry()
modes_state, modes_config = make_mode_response(
list(agent_loop.agent_manager.available_agents.values()),
session.agent_loop.agent_profile.name,
)
models_state, models_config = make_model_response(
agent_loop.config.models, agent_loop.config.active_model
)
modes_state, _, models_state, _ = self._build_session_state(session)
return NewSessionResponse(
session_id=session.id,
models=models_state,
modes=modes_state,
config_options=[modes_config, models_config],
config_options=self._build_config_options(session),
)
def _get_acp_tool_overrides(self) -> list[Path]:
@ -439,7 +482,7 @@ class VibeAcpAgentLoop(AcpAgent):
except Exception:
pass
asyncio.create_task(_send())
session.spawn(_send())
async def _replay_tool_calls(self, session_id: str, msg: LLMMessage) -> None:
if not msg.tool_calls:
@ -476,9 +519,7 @@ class VibeAcpAgentLoop(AcpAgent):
session_id=session_id, update=result_update
)
async def _send_available_commands(self, session_id: str) -> None:
session = self._get_session(session_id)
async def _send_available_commands(self, session: AcpSessionLoop) -> None:
commands: list[AvailableCommand] = []
for cmd in session.command_registry.commands.values():
@ -510,7 +551,7 @@ class VibeAcpAgentLoop(AcpAgent):
)
await self.client.session_update(
session_id=session_id, update=update_available_commands(commands)
session_id=session.id, update=update_available_commands(commands)
)
@override
@ -525,6 +566,7 @@ class VibeAcpAgentLoop(AcpAgent):
os.chdir(cwd)
config = self._load_config()
hook_config_result = load_hooks_from_fs(config)
session_dir = SessionLoader.find_session_by_id(
session_id, config.session_logging
@ -533,42 +575,35 @@ class VibeAcpAgentLoop(AcpAgent):
raise SessionNotFoundError(session_id)
try:
loaded_messages, _ = SessionLoader.load_session(session_dir)
loaded_messages, metadata = SessionLoader.load_session(session_dir)
except Exception as e:
raise SessionLoadError(session_id, str(e)) from e
agent_loop = AgentLoop(
config=config,
agent_name=BuiltinAgentName.DEFAULT,
enable_streaming=True,
entrypoint_metadata=self._build_entrypoint_metadata(),
defer_heavy_init=True,
agent_loop = self._create_agent_loop(
config, BuiltinAgentName.DEFAULT, hook_config_result=hook_config_result
)
loaded_session_id = metadata.get("session_id", agent_loop.session_id)
agent_loop.session_id = loaded_session_id
agent_loop.parent_session_id = metadata.get("parent_session_id")
agent_loop.session_logger.resume_existing_session(
loaded_session_id, session_dir
)
agent_loop.agent_manager.register_agent(CHAT_AGENT)
non_system_messages = [
msg for msg in loaded_messages if msg.role != Role.system
]
agent_loop.messages.extend(non_system_messages)
if non_system_messages:
agent_loop.messages.extend(non_system_messages)
session = await self._create_acp_session(session_id, agent_loop)
await self._replay_conversation_history(session_id, non_system_messages)
await self._replay_conversation_history(session.id, non_system_messages)
self._send_usage_update(session)
modes_state, modes_config = make_mode_response(
list(agent_loop.agent_manager.available_agents.values()),
session.agent_loop.agent_profile.name,
)
models_state, models_config = make_model_response(
agent_loop.config.models, agent_loop.config.active_model
)
modes_state, _, models_state, _ = self._build_session_state(session)
return LoadSessionResponse(
models=models_state,
modes=modes_state,
config_options=[modes_config, models_config],
config_options=self._build_config_options(session),
)
async def _apply_mode_change(self, session: AcpSessionLoop, mode_id: str) -> bool:
@ -578,7 +613,7 @@ class VibeAcpAgentLoop(AcpAgent):
await session.agent_loop.switch_agent(mode_id)
if session.agent_loop.auto_approve:
if session.agent_loop.bypass_tool_permissions:
session.agent_loop.approval_callback = None
else:
session.agent_loop.set_approval_callback(
@ -587,20 +622,27 @@ class VibeAcpAgentLoop(AcpAgent):
return True
async def _reload_config(self, session: AcpSessionLoop) -> None:
new_config = VibeConfig.load(
tool_paths=session.agent_loop.config.tool_paths,
disabled_tools=["ask_user_question"],
)
await session.agent_loop.reload_with_initial_messages(base_config=new_config)
async def _apply_model_change(self, session: AcpSessionLoop, model_id: str) -> bool:
model_aliases = [model.alias for model in session.agent_loop.config.models]
if model_id not in model_aliases:
return False
VibeConfig.save_updates({"active_model": model_id})
await self._reload_config(session)
return True
new_config = VibeConfig.load(
tool_paths=session.agent_loop.config.tool_paths,
disabled_tools=["ask_user_question"],
)
await session.agent_loop.reload_with_initial_messages(base_config=new_config)
async def _apply_thinking_change(
self, session: AcpSessionLoop, level: ThinkingLevel
) -> bool:
session.agent_loop.config.set_thinking(level)
await self._reload_config(session)
return True
@override
@ -636,6 +678,10 @@ class VibeAcpAgentLoop(AcpAgent):
success = await self._apply_mode_change(session, value)
case "model" if isinstance(value, str):
success = await self._apply_model_change(session, value)
case "thinking" if isinstance(value, str) and value in THINKING_LEVELS:
success = await self._apply_thinking_change(
session, cast(ThinkingLevel, value)
)
case _:
success = False
@ -682,7 +728,7 @@ class VibeAcpAgentLoop(AcpAgent):
) -> PromptResponse:
session = self._get_session(session_id)
if session.task is not None:
if session.prompt_task is not None:
raise InvalidRequestError(
"Concurrent prompts are not supported yet, wait for agent loop to finish"
)
@ -713,8 +759,8 @@ class VibeAcpAgentLoop(AcpAgent):
await self.client.session_update(session_id=session.id, update=update)
try:
session.task = asyncio.create_task(agent_loop_task())
await session.task
task = session.set_prompt_task(agent_loop_task())
await task
except asyncio.CancelledError:
self._send_usage_update(session)
@ -727,15 +773,15 @@ class VibeAcpAgentLoop(AcpAgent):
except CoreRateLimitError as e:
raise RateLimitError.from_core(e) from e
except CoreContextTooLongError as e:
raise ContextTooLongError.from_core(e) from e
except ConversationLimitException as e:
raise ConversationLimitError(str(e)) from e
except Exception as e:
raise InternalError(str(e)) from e
finally:
session.task = None
self._send_usage_update(session)
return PromptResponse(
stop_reason="end_turn",
@ -803,12 +849,13 @@ class VibeAcpAgentLoop(AcpAgent):
if command is None:
return None
session.agent_loop.telemetry_client.send_slash_command_used(cmd_name, "builtin")
handler = getattr(self, command.handler)
return await handler(session, text_prompt, message_id)
async def _run_agent_loop(
self, session: AcpSessionLoop, prompt: str, client_message_id: str | None = None
) -> AsyncGenerator[SessionUpdate]:
) -> AsyncGenerator[SessionUpdate | UsageUpdate]:
rendered_prompt = render_path_prompt(prompt, base_dir=Path.cwd())
async with aclosing(
@ -846,6 +893,7 @@ class VibeAcpAgentLoop(AcpAgent):
session_update = tool_result_session_update(event)
if session_update:
yield session_update
self._send_usage_update(session)
elif isinstance(event, ToolStreamEvent):
yield ToolCallProgress(
@ -874,14 +922,31 @@ class VibeAcpAgentLoop(AcpAgent):
async def close_session(
self, session_id: str, **kwargs: Any
) -> CloseSessionResponse | None:
raise NotImplementedMethodError("close_session")
session = self._get_session(session_id)
self.sessions.pop(session_id, None)
await session.close()
await self._close_agent_loop(session.agent_loop)
return CloseSessionResponse()
async def _close_agent_loop(self, agent_loop: AgentLoop) -> None:
deferred_init_thread = agent_loop._deferred_init_thread
if deferred_init_thread is not None and deferred_init_thread.is_alive():
await asyncio.to_thread(deferred_init_thread.join)
backend_close = getattr(agent_loop.backend, "close", None)
if callable(backend_close):
close_result = backend_close()
if inspect.isawaitable(close_result):
await close_result
await agent_loop.telemetry_client.aclose()
@override
async def cancel(self, session_id: str, **kwargs: Any) -> None:
session = self._get_session(session_id)
if session.task and not session.task.done():
session.task.cancel()
session.task = None
await session.cancel_prompt()
@override
async def fork_session(
@ -891,7 +956,41 @@ class VibeAcpAgentLoop(AcpAgent):
mcp_servers: list[HttpMcpServer | SseMcpServer | McpServerStdio] | None = None,
**kwargs: Any,
) -> ForkSessionResponse:
raise NotImplementedMethodError("fork_session")
load_dotenv_values()
os.chdir(cwd)
source_session = self._get_session(session_id)
try:
message_id = ForkSessionParams.model_validate(kwargs).message_id
except ValidationError as e:
raise InvalidRequestError(f"Invalid fork parameters: {e}") from e
if (
source_session.prompt_task is not None
and not source_session.prompt_task.done()
):
raise InvalidRequestError(
"Cannot fork a session while the agent loop is running"
)
try:
agent_loop = await source_session.agent_loop.fork(message_id)
agent_loop.agent_manager.register_agent(CHAT_AGENT)
session = await self._create_acp_session(agent_loop.session_id, agent_loop)
except InvalidRequestError:
raise
except ValueError as e:
raise InvalidRequestError(str(e)) from e
except Exception as e:
raise ConfigurationError(str(e)) from e
modes_state, _, models_state, _ = self._build_session_state(session)
return ForkSessionResponse(
session_id=session.id,
models=models_state,
modes=modes_state,
config_options=self._build_config_options(session),
)
@override
async def resume_session(
@ -909,7 +1008,34 @@ class VibeAcpAgentLoop(AcpAgent):
@override
async def ext_notification(self, method: str, params: dict) -> None:
raise NotImplementedMethodError("ext_notification")
# ACP strips the leading "_" before delegating extension notifications here.
if method == "telemetry/send":
self._handle_telemetry_notification(params)
def _handle_telemetry_notification(self, params: dict[str, Any]) -> None:
try:
notification = TelemetrySendNotification.model_validate(params)
except ValidationError as exc:
raise InvalidRequestError(
f"Invalid ACP telemetry notification: {exc}"
) from exc
session = self.sessions.get(notification.session_id)
if session is None:
logger.warning(
"Ignoring ACP telemetry notification because session could not be resolved: %s",
notification.session_id,
)
return
dispatcher = _EVENT_DISPATCHERS.get(notification.event)
if dispatcher is None:
logger.warning(
"Ignoring unsupported ACP telemetry event: %s", notification.event
)
return
dispatcher(session.agent_loop.telemetry_client, notification.properties)
@override
def on_connect(self, conn: Client) -> None:
@ -962,6 +1088,8 @@ class VibeAcpAgentLoop(AcpAgent):
tool_call_id = str(uuid4())
old_tokens = session.agent_loop.stats.context_tokens
parts = text_prompt.strip().split(None, 1)
cmd_args = parts[1] if len(parts) > 1 else ""
start_event = CompactStartEvent(
current_context_tokens=old_tokens or 0,
@ -973,7 +1101,7 @@ class VibeAcpAgentLoop(AcpAgent):
update=create_compact_start_session_update(start_event),
)
await session.agent_loop.compact()
await session.agent_loop.compact(extra_instructions=cmd_args.strip())
new_tokens = session.agent_loop.stats.context_tokens
end_event = CompactEndEvent(
@ -1070,13 +1198,16 @@ class VibeAcpAgentLoop(AcpAgent):
) -> list[SessionConfigOptionSelect | SessionConfigOptionBoolean]:
"""Build the current modes + models config options for a session."""
profiles = list(session.agent_loop.agent_manager.available_agents.values())
_, modes_config = make_mode_response(
_, modes_config = build_mode_state(
profiles, session.agent_loop.agent_profile.name
)
_, models_config = make_model_response(
_, models_config = build_model_state(
session.agent_loop.config.models, session.agent_loop.config.active_model
)
return [modes_config, models_config]
thinking_config = make_thinking_response(
session.agent_loop.config.get_active_model().thinking
)
return [modes_config, models_config, thinking_config]
async def _send_config_option_update(self, session: AcpSessionLoop) -> None:
"""Push updated config options (modes, models) to the client."""