diff --git a/.gitignore b/.gitignore index c6cf4399..9739f3e8 100644 --- a/.gitignore +++ b/.gitignore @@ -25,6 +25,7 @@ Thumbs.db *.sqlite3 .coverage .htmlcov/ +.firecrawl/ *.tmp *.temp basilisk_config.yml diff --git a/README.md b/README.md index b86e5e06..f8a383b9 100644 --- a/README.md +++ b/README.md @@ -71,6 +71,8 @@ You can obtain an API key from your chosen AI provider. The app supports the fol - [Mistral AI](https://mistral.ai/) - [Ollama](https://ollama.com/) - [OpenAI](https://platform.openai.com/) +- [OpenCode Go](https://opencode.ai/docs/go/) +- [OpenCode Zen](https://opencode.ai/docs/zen/) - [OpenRouter](https://openrouter.ai/) - [xAI](https://x.ai/) diff --git a/basilisk/completion_handler.py b/basilisk/completion_handler.py index b49dd4aa..c926bb54 100644 --- a/basilisk/completion_handler.py +++ b/basilisk/completion_handler.py @@ -32,6 +32,9 @@ logger = logging.getLogger(__name__) +# Upper bound for joining a worker on the UI thread after a failed start. +_START_FAILURE_JOIN_TIMEOUT = 1.0 + COMMON_PATTERN = r"[\n;:.?!)ยป\"\]}]" RE_STREAM_BUFFER = re.compile(rf".*{COMMON_PATTERN}.*") @@ -77,8 +80,14 @@ def __init__( self.on_non_stream_finish = on_non_stream_finish self.task: Optional[threading.Thread] = None self._stop_completion = False + self._completion_lock = threading.Lock() + self._active_request: object | None = None + self._latest_request: object | None = None + self._active_engine: Optional[BaseEngine] = None + self._active_response: Any = None + self._start_notified: threading.Event | None = None self.last_time = 0 - self.stream_buffer: str = "" + self._stream_buffers: dict[object, str] = {} @ensure_no_task_running def start_completion( @@ -100,8 +109,7 @@ def start_completion( stream: Whether to use streaming mode **kwargs: Additional arguments for the completion """ - self._stop_completion = False - + request = object() completion_args = { "engine": engine, "system_message": system_message, @@ -110,15 +118,63 @@ def start_completion( "stream": stream, **kwargs, } + start_notified = threading.Event() + with self._completion_lock: + self._stop_completion = False + self._active_request = request + self._latest_request = request + self._active_engine = None + self._active_response = None + self._start_notified = start_notified + self._stream_buffers[request] = "" + task = threading.Thread( + target=self._handle_completion, + kwargs={ + "request": request, + "start_notified": start_notified, + **completion_args, + }, + ) + self.task = task + try: + task.start() + except Exception: + self.task = None + self._active_request = None + self._latest_request = None + self._active_engine = None + self._active_response = None + self._start_notified = None + self._stream_buffers.pop(request, None) + raise - if self.on_completion_start: - self.on_completion_start() - - self.task = threading.Thread( - target=self._handle_completion, kwargs=completion_args - ) - self.task.start() - logger.debug("Completion task %s started", self.task.ident) + try: + if self.on_completion_start: + self.on_completion_start() + except Exception: + task_to_join = None + to_cancel = None + with self._completion_lock: + if self._active_request is request: + self._stop_completion = True + task_to_join = self.task + to_cancel = self._cancellable_response() + start_notified.set() + if to_cancel is not None: + self._cancel_response(*to_cancel) + if task_to_join is not None: + # This runs on the UI thread. The worker is at most waiting for + # start_notified and sees the stop flag at once, so the join is + # normally immediate; the bound keeps a wedged provider from + # freezing the UI, and the worker clears self.task itself. + task_to_join.join(timeout=_START_FAILURE_JOIN_TIMEOUT) + with self._completion_lock: + if self._latest_request is request: + self._latest_request = None + raise + else: + start_notified.set() + logger.debug("Completion task %s started", task.ident) def stop_completion(self, skip_callbacks: bool = False): """Stop the current completion if running. @@ -127,78 +183,258 @@ def stop_completion(self, skip_callbacks: bool = False): skip_callbacks: If True, skip calling completion end callbacks. Useful when cleaning up resources before destroying the tab. """ - if self.is_running(): - self._stop_completion = True - logger.debug("Stopping completion task: %s", self.task.ident) - self.task.join(timeout=0.05) - self.task = None + start_notified = None + to_cancel = None + with self._completion_lock: + task = self.task + request = self._latest_request + is_running = bool(task and task.is_alive()) + is_active_request = ( + request is not None and self._active_request is request + ) + if is_running or is_active_request: + self._stop_completion = True + to_cancel = self._cancellable_response() + start_notified = self._start_notified + if skip_callbacks and self._latest_request is request: + self._latest_request = None + if not is_running and is_active_request: + self.task = None + self._active_request = None + self._active_engine = None + self._active_response = None + self._start_notified = None + if request is not None: + self._stream_buffers.pop(request, None) + + # Both are None unless this call is stopping a request. + if start_notified is not None: + start_notified.set() + if to_cancel is not None: + self._cancel_response(*to_cancel) + if is_running: + logger.debug("Stopping completion task: %s", task.ident) + task.join(timeout=0.1) + with self._completion_lock: + if self.task is task and not task.is_alive(): + self.task = None + if is_running or is_active_request: + stop_sound() if self.on_completion_end and not skip_callbacks: - wx.CallAfter(self.on_completion_end, False) + if request is None: + wx.CallAfter(self.on_completion_end, False) + else: + wx.CallAfter(self._completion_stopped, request) def is_running(self) -> bool: """Check if a completion is currently running.""" - return self.task and self.task.is_alive() + with self._completion_lock: + return bool(self.task and self.task.is_alive()) + + def _cancellable_response(self) -> tuple[BaseEngine, Any] | None: + """Return the active engine and response to cancel; hold the lock.""" + if self._active_engine is None or self._active_response is None: + return None + return self._active_engine, self._active_response + + def _cancel_response(self, engine: BaseEngine, response: Any) -> None: + """Close an active provider response without surfacing stop errors.""" + try: + engine.cancel_completion(response) + except Exception: + logger.debug( + "Provider response could not be closed cleanly", exc_info=True + ) + + def _is_stopping(self, request: object) -> bool: + with self._completion_lock: + return ( + self._stop_completion + or self._active_request is not request + or global_vars.app_should_exit + ) + + def _wait_for_start_notification( + self, request: object, start_notified: threading.Event | None + ) -> bool: + """Wait for the start callback and confirm this request still owns work.""" + if start_notified is not None: + start_notified.wait() + return not self._is_stopping(request) + + def _queue_error_if_active( + self, request: object, error_message: str + ) -> None: + """Queue an error only while its request still owns the lifecycle.""" + if not self._is_stopping(request): + wx.CallAfter(self._handle_error, request, error_message) + + def _commit_startup(self, request: object) -> bool: + """Start progress sound while atomically committing request startup.""" + with self._completion_lock: + if ( + self._stop_completion + or self._active_request is not request + or global_vars.app_should_exit + ): + return False + play_sound("progress", loop=True) + return True + + def _start_provider_completion( + self, + request: object, + engine: BaseEngine, + completion_args: dict[str, Any], + ) -> tuple[bool, Any]: + """Commit startup and obtain a provider response for an active request.""" + try: + if not self._commit_startup(request): + return False, None + return True, engine.completion(**completion_args) + except Exception as e: + if self._is_stopping(request): + logger.debug("Completion request cancelled", exc_info=True) + else: + logger.error("Error during completion", exc_info=True) + self._queue_error_if_active(request, str(e)) + return False, None - def _handle_completion(self, engine: BaseEngine, **kwargs: dict[str, Any]): + def _publish_response( + self, request: object, engine: BaseEngine, response: Any + ) -> bool: + """Publish a response and report whether its request was stopped.""" + with self._completion_lock: + is_active_request = self._active_request is request + if is_active_request: + self._active_engine = engine + self._active_response = response + return ( + self._stop_completion + or global_vars.app_should_exit + or not is_active_request + ) + + def _handle_completion( + self, + request: object, + engine: BaseEngine, + start_notified: threading.Event | None = None, + **kwargs: dict[str, Any], + ): """Handle the completion request in a background thread. Args: + request: Identity of the worker handling this completion. engine: The engine to use for completion + start_notified: Set after the completion-start callback returns. kwargs: The keyword arguments for the completion request """ try: - play_sound("progress", loop=True) - response = engine.completion(**kwargs) - except Exception as e: - logger.error("Error during completion", exc_info=True) - wx.CallAfter(self._handle_error, str(e)) - return + if not self._wait_for_start_notification(request, start_notified): + return + started, response = self._start_provider_completion( + request, engine, kwargs + ) + if not started: + return - handle_func = ( - self._handle_streaming_completion - if kwargs.get("stream", False) - else self._handle_non_streaming_completion - ) - kwargs["engine"] = engine - kwargs["response"] = response - try: - success = handle_func(**kwargs) - except Exception as e: - logger.error("Error handling completion response", exc_info=True) - wx.CallAfter(self._handle_error, str(e)) - return + if self._publish_response(request, engine, response): + self._cancel_response(engine, response) + return - if success: - wx.CallAfter(self._completion_finished_success) + handle_func = ( + self._handle_streaming_completion + if kwargs.get("stream", False) + else self._handle_non_streaming_completion + ) + kwargs["engine"] = engine + kwargs["response"] = response + try: + success = handle_func(request=request, **kwargs) + except Exception as e: + if self._is_stopping(request): + logger.debug("Completion response cancelled", exc_info=True) + return + logger.error( + "Error handling completion response", exc_info=True + ) + self._queue_error_if_active(request, str(e)) + return + + if success and not self._is_stopping(request): + wx.CallAfter(self._completion_finished_success, request) + finally: + with self._completion_lock: + self._stream_buffers.pop(request, None) + if self._active_request is request: + self._active_engine = None + self._active_response = None + if ( + self._stop_completion + and self.task is threading.current_thread() + ): + self.task = None + if self._active_request is request: + self._active_request = None + if self._start_notified is start_notified: + self._start_notified = None def _handle_stream_chunk( - self, chunk: str | tuple[str, Any], message_block: MessageBlock + self, + request: object, + chunk: str | tuple[str, Any], + message_block: MessageBlock, ): - if isinstance(chunk, str): - self.stream_buffer += chunk - elif isinstance(chunk, tuple): - chunk_type, chunk_data = chunk - if chunk_type == "citation": - if not message_block.response.citations: - message_block.response.citations = [] - message_block.response.citations.append(chunk_data) + with self._completion_lock: + if ( + self._stop_completion + or self._active_request is not request + or global_vars.app_should_exit + ): + return + if isinstance(chunk, str): + buffer = self._stream_buffers.get(request, "") + chunk + self._stream_buffers[request] = buffer + elif isinstance(chunk, tuple): + chunk_type, chunk_data = chunk + if chunk_type == "citation": + if not message_block.response.citations: + message_block.response.citations = [] + message_block.response.citations.append(chunk_data) + else: + logger.warning( + "Unknown chunk type in streaming response: %s", + chunk_type, + ) + buffer = self._stream_buffers.get(request, "") else: - logger.warning( - "Unknown chunk type in streaming response: %s", chunk_type - ) + buffer = self._stream_buffers.get(request, "") - if RE_STREAM_BUFFER.match(self.stream_buffer): - self.flush_stream_buffer(message_block) + if RE_STREAM_BUFFER.match(buffer): + self.flush_stream_buffer(request, message_block) - def flush_stream_buffer(self, message_block: MessageBlock) -> None: + def flush_stream_buffer( + self, request: object, message_block: MessageBlock + ) -> None: """Flush the stream buffer to the message block.""" - if self.stream_buffer: - message_block.response.content += self.stream_buffer - wx.CallAfter(self._handle_stream_buffer, self.stream_buffer) - self.stream_buffer = "" + with self._completion_lock: + if ( + self._stop_completion + or self._active_request is not request + or global_vars.app_should_exit + ): + return + buffer = self._stream_buffers.get(request, "") + self._stream_buffers[request] = "" + + if buffer: + message_block.response.content += buffer + wx.CallAfter(self._handle_stream_buffer, request, buffer) def _handle_streaming_completion( self, + request: object, engine: BaseEngine, response: Any, new_block: MessageBlock, @@ -208,6 +444,7 @@ def _handle_streaming_completion( """Handle streaming completion response. Args: + request: Identity of the worker handling this completion. engine: The engine used for completion response: The completion response new_block: The message block being completed @@ -221,22 +458,25 @@ def _handle_streaming_completion( # Notify that streaming has started if self.on_stream_start: - wx.CallAfter(self.on_stream_start, new_block, system_message) + wx.CallAfter( + self._handle_stream_start, request, new_block, system_message + ) for chunk in engine.completion_response_with_stream(response): - if self._stop_completion or global_vars.app_should_exit: + if self._is_stopping(request): logger.debug("Stopping completion") return False - self._handle_stream_chunk(chunk, new_block) + self._handle_stream_chunk(request, chunk, new_block) # Notify that streaming has finished - self.flush_stream_buffer(new_block) + self.flush_stream_buffer(request, new_block) if self.on_stream_finish: - wx.CallAfter(self.on_stream_finish, new_block) + wx.CallAfter(self._handle_stream_finish, request, new_block) return True def _handle_non_streaming_completion( self, + request: object, engine: BaseEngine, response: Any, new_block: MessageBlock, @@ -246,6 +486,7 @@ def _handle_non_streaming_completion( """Handle non-streaming completion response. Args: + request: Identity of the worker handling this completion. engine: The engine used for completion response: The completion response new_block: The message block being completed @@ -262,17 +503,55 @@ def _handle_non_streaming_completion( # Notify that non-streaming completion has finished if self.on_non_stream_finish: wx.CallAfter( - self.on_non_stream_finish, completed_block, system_message + self._handle_non_stream_finish, + request, + completed_block, + system_message, ) return True - def _handle_stream_buffer(self, buffer: str): + def _handle_stream_start( + self, + request: object, + new_block: MessageBlock, + system_message: Optional[SystemMessage], + ): + """Notify the UI that the owning stream has started.""" + if self._is_stopping(request): + return + if self.on_stream_start: + self.on_stream_start(new_block, system_message) + + def _handle_stream_finish(self, request: object, new_block: MessageBlock): + """Notify the UI that the owning stream has finished.""" + if self._is_stopping(request): + return + if self.on_stream_finish: + self.on_stream_finish(new_block) + + def _handle_non_stream_finish( + self, + request: object, + completed_block: MessageBlock, + system_message: Optional[SystemMessage], + ): + """Notify the UI that the owning non-streaming request finished.""" + if self._is_stopping(request): + return + if self.on_non_stream_finish: + self.on_non_stream_finish(completed_block, system_message) + + def _handle_stream_buffer(self, request: object, buffer: str): """Handle a streaming chunk on the main thread. Args: + request: Identity of the worker that produced the buffer. buffer: The streaming buffer content """ + if self._is_stopping(request): + return + if self.on_stream_chunk: self.on_stream_chunk(buffer) @@ -282,20 +561,39 @@ def _handle_stream_buffer(self, buffer: str): play_sound("chat_response_pending") self.last_time = new_time - def _completion_finished_success(self): + def _completion_finished_success(self, request: object): """Handle completion finish in success on the main thread.""" + with self._completion_lock: + if self._stop_completion or self._active_request is not request: + return + self.task = None + self._active_request = None + self._active_engine = None + self._active_response = None + self._start_notified = None + self._stream_buffers.pop(request, None) stop_sound() play_sound("chat_response_received") if self.on_completion_end: self.on_completion_end(True) - self.task = None - def _handle_error(self, error_message: str): + def _handle_error(self, request: object, error_message: str): """Handle completion error on the main thread. Args: + request: Identity of the worker that raised the error. error_message: The error message """ + with self._completion_lock: + if self._stop_completion or self._active_request is not request: + return + self.task = None + self._active_request = None + self._active_engine = None + self._active_response = None + self._start_notified = None + self._stream_buffers.pop(request, None) + stop_sound() if self.on_error: @@ -312,4 +610,10 @@ def _handle_error(self, error_message: str): if self.on_completion_end: self.on_completion_end(False) - self.task = None + def _completion_stopped(self, request: object): + """Notify the UI that the latest completion was stopped.""" + with self._completion_lock: + if self._latest_request is not request: + return + if self.on_completion_end: + self.on_completion_end(False) diff --git a/basilisk/presenters/conversation_presenter.py b/basilisk/presenters/conversation_presenter.py index cda6e7fa..3713bb89 100644 --- a/basilisk/presenters/conversation_presenter.py +++ b/basilisk/presenters/conversation_presenter.py @@ -197,9 +197,8 @@ def get_completion_args(self) -> dict[str, Any] | None: def cleanup(self): """Stop all active resources before destroying the tab.""" - if self.completion_handler.is_running(): - log.debug("Stopping completion handler before closing tab") - self.completion_handler.stop_completion(skip_callbacks=True) + log.debug("Stopping completion handler before closing tab") + self.completion_handler.stop_completion(skip_callbacks=True) if self.recording_thread and self.recording_thread.is_alive(): log.debug("Aborting recording thread before closing tab") try: diff --git a/basilisk/provider.py b/basilisk/provider.py index 61f7be91..e859a85c 100644 --- a/basilisk/provider.py +++ b/basilisk/provider.py @@ -139,6 +139,30 @@ def engine_cls(self) -> Type[BaseEngine]: engine_cls_path="basilisk.provider_engine.openai_engine.OpenAIEngine", allow_custom_base_url=True, ), + Provider( + id="opencodego", + name="OpenCode Go", + base_url="https://opencode.ai/zen/go/v1", + api_type=ProviderAPIType.OPENAI, + organization_mode_available=False, + require_api_key=True, + env_var_name_api_key="OPENCODE_API_KEY", + engine_cls_path=( + "basilisk.provider_engine.opencode_go_engine.OpenCodeGoEngine" + ), + ), + Provider( + id="opencodezen", + name="OpenCode Zen", + base_url="https://opencode.ai/zen/v1", + api_type=ProviderAPIType.OPENAI, + organization_mode_available=False, + require_api_key=True, + env_var_name_api_key="OPENCODE_API_KEY", + engine_cls_path=( + "basilisk.provider_engine.opencode_go_engine.OpenCodeZenEngine" + ), + ), Provider( id="openrouter", name="OpenRouter", diff --git a/basilisk/provider_engine/base_engine.py b/basilisk/provider_engine/base_engine.py index df707ba9..a02f32e2 100644 --- a/basilisk/provider_engine/base_engine.py +++ b/basilisk/provider_engine/base_engine.py @@ -419,6 +419,17 @@ def completion_response_with_stream(self, stream: Any, **kwargs) -> Any: """ pass + def cancel_completion(self, response: Any) -> None: + """Cancel an active completion response when the SDK supports it. + + Most synchronous provider SDKs expose streaming responses with a + ``close`` method. Engines with a different cancellation mechanism can + override this hook. + """ + close = getattr(response, "close", None) + if callable(close): + close() + @abstractmethod def completion_response_without_stream( self, response: Any, new_block: MessageBlock, **kwargs diff --git a/basilisk/provider_engine/opencode_go_engine.py b/basilisk/provider_engine/opencode_go_engine.py new file mode 100644 index 00000000..5d3eab1d --- /dev/null +++ b/basilisk/provider_engine/opencode_go_engine.py @@ -0,0 +1,411 @@ +"""OpenCode Go subscription API integration.""" + +from __future__ import annotations + +import enum +import logging +from dataclasses import dataclass +from functools import cached_property +from typing import TYPE_CHECKING, Any, ClassVar, Iterator, Mapping + +import httpx +from anthropic import Anthropic +from google import genai +from google.genai.types import HttpOptions, ResourceScope + +from basilisk.conversation import Conversation, Message, MessageBlock +from basilisk.provider_ai_model import ProviderAIModel +from basilisk.provider_capability import ProviderCapability + +from .anthropic_engine import AnthropicEngine +from .gemini_engine import GeminiEngine +from .legacy_openai_engine import LegacyOpenAIEngine +from .openai_engine import OpenAIEngine + +if TYPE_CHECKING: + from basilisk.config import Account + +log = logging.getLogger(__name__) + + +class _Protocol(enum.Enum): + CHAT_COMPLETIONS = enum.auto() + RESPONSES = enum.auto() + ANTHROPIC_MESSAGES = enum.auto() + GEMINI = enum.auto() + + +@dataclass(frozen=True) +class _ProtocolResponse: + protocol: _Protocol + value: Any + + +# OpenCode Go exposes models through different API protocols. Keep these maps +# aligned with https://opencode.ai/docs/go/#endpoints. Newly added models use +# chat completions until their protocol is documented here. +_RESPONSES_MODELS = frozenset( + {"grok-4.5", "gpt-5.6-luna", "muse-spark-1.2-contributor"} +) +_ANTHROPIC_MODELS = frozenset( + { + "minimax-m3", + "minimax-m2.7", + "minimax-m2.5", + "qwen3.8-max", + "qwen3.7-max", + "qwen3.7-plus", + "qwen3.6-plus", + } +) +_GO_VISION_MODELS = frozenset( + { + "deepseek-v4-flash-vision-exp", + "gpt-5.6-luna", + "grok-4.5", + "kimi-k2.5", + "kimi-k2.6", + "kimi-k3", + "kimi-k2.7-code", + "mimo-v2-omni", + "mimo-v2-pro", + "mimo-v2.5", + "mimo-v2.5-pro", + "muse-spark-1.2-contributor", + "ox-alpha-free", + "qwen3.5-plus", + "qwen3.6-plus", + "qwen3.7-plus", + "qwen3.8-max", + } +) +_GO_FIXED_KIMI_SAMPLING_PARAMETERS = { + "kimi-k2.5": {"temperature": 1.0, "top_p": 0.95}, + "kimi-k2.6": {"temperature": 1.0, "top_p": 0.95}, + "kimi-k2.7-code": {"temperature": 1.0, "top_p": 0.95}, + "kimi-k3": {"temperature": 1.0, "top_p": 0.95}, +} +_UNSUPPORTED_SAMPLING_PARAMETERS = ("temperature", "top_p") +_GO_MODEL_UNSUPPORTED_PARAMETERS = { + model_id: _UNSUPPORTED_SAMPLING_PARAMETERS + for model_id in (*_GO_FIXED_KIMI_SAMPLING_PARAMETERS, "gpt-5.6-luna") +} +_ZEN_MODEL_UNSUPPORTED_PARAMETERS = { + "kimi-k3": _UNSUPPORTED_SAMPLING_PARAMETERS +} + + +class _SharedModelsMixin: + """Use the parent OpenCode engine's discovered model list.""" + + def __init__(self, account: Account, model_source: _OpenCodeEngine): + self._model_source = model_source + super().__init__(account) + + @property + def models(self) -> list[ProviderAIModel]: + return self._model_source.models + + +class _OpenCodeResponsesEngine(_SharedModelsMixin, OpenAIEngine): + """OpenAI Responses adapter sharing the OpenCode model catalog.""" + + +class _OpenCodeAnthropicEngine(_SharedModelsMixin, AnthropicEngine): + """Anthropic Messages adapter sharing the OpenCode model catalog.""" + + @cached_property + def client(self) -> Anthropic: + base_url = str(self.account.provider.base_url).rstrip("/") + return Anthropic( + api_key=self.account.api_key.get_secret_value(), + base_url=base_url.removesuffix("/v1"), + ) + + +class _OpenCodeGeminiEngine(_SharedModelsMixin, GeminiEngine): + """Gemini adapter sharing the OpenCode model catalog.""" + + @cached_property + def client(self) -> genai.Client: + return genai.Client( + api_key=self.account.api_key.get_secret_value(), + http_options=HttpOptions( + base_url=str(self.account.provider.base_url).rstrip("/"), + base_url_resource_scope=ResourceScope.COLLECTION, + ), + ) + + +class _OpenCodeEngine(LegacyOpenAIEngine): + """Route OpenCode models to their documented API protocol.""" + + capabilities: set[ProviderCapability] = { + ProviderCapability.IMAGE, + ProviderCapability.TEXT, + } + supported_attachment_formats: set[str] = { + "image/gif", + "image/jpeg", + "image/png", + "image/webp", + } + VISION_MODELS: ClassVar[frozenset[str]] = frozenset() + FIXED_SAMPLING_PARAMETERS: ClassVar[Mapping[str, Mapping[str, float]]] = {} + MODEL_UNSUPPORTED_PARAMETERS: ClassVar[Mapping[str, tuple[str, ...]]] = {} + + @cached_property + def _responses_engine(self) -> _OpenCodeResponsesEngine: + return _OpenCodeResponsesEngine(self.account, self) + + @cached_property + def _anthropic_engine(self) -> _OpenCodeAnthropicEngine: + return _OpenCodeAnthropicEngine(self.account, self) + + @cached_property + def _gemini_engine(self) -> _OpenCodeGeminiEngine: + return _OpenCodeGeminiEngine(self.account, self) + + def _load_models(self) -> list[ProviderAIModel]: + """Fetch the models currently available from this provider.""" + url = f"{str(self.account.provider.base_url).rstrip('/')}/models" + response = httpx.get( + url, + headers={ + "Authorization": ( + f"Bearer {self.account.api_key.get_secret_value()}" + ), + "User-Agent": self.get_user_agent(), + }, + timeout=30.0, + ) + if response.status_code != 200: + raise httpx.HTTPStatusError( + ( + f"Failed to get models from '{url}' " + f"(status={response.status_code}): {response.text}" + ), + request=response.request, + response=response, + ) + data = response.json() + rows = data.get("data", []) + if not isinstance(rows, list): + raise ValueError("OpenCode returned an invalid model list") + + models = [] + for row in rows: + if not isinstance(row, dict): + continue + model_id = row.get("id") + if not isinstance(model_id, str) or not model_id: + continue + created = row.get("created", 0) + if not isinstance(created, int): + created = 0 + extra_info = {"owned_by": row.get("owned_by")} + if unsupported_parameters := self.MODEL_UNSUPPORTED_PARAMETERS.get( + model_id + ): + extra_info["unsupported_parameters"] = list( + unsupported_parameters + ) + models.append( + ProviderAIModel( + id=model_id, + # Translators: Description for a model dynamically discovered + # from an OpenCode provider. + description=_( + "Model available from this OpenCode provider." + ), + max_output_tokens=8192, + vision=model_id in self.VISION_MODELS, + created=created, + extra_info=extra_info, + ) + ) + return models + + @staticmethod + def _protocol_for_model(model_id: str) -> _Protocol: + if model_id in _RESPONSES_MODELS: + return _Protocol.RESPONSES + if model_id in _ANTHROPIC_MODELS: + return _Protocol.ANTHROPIC_MESSAGES + return _Protocol.CHAT_COMPLETIONS + + def _strip_catalog_sampling_params( + self, model: ProviderAIModel | None, params: dict[str, Any] + ) -> None: + """Apply the Go-only fixed Kimi sampling policy after catalog stripping.""" + super()._strip_catalog_sampling_params(model, params) + if model is not None: + params.update(self.FIXED_SAMPLING_PARAMETERS.get(model.id, {})) + + def completion( + self, + new_block: MessageBlock, + conversation: Conversation, + system_message: Message | None, + stop_block_index: int | None = None, + **kwargs, + ) -> _ProtocolResponse: + """Create a completion through the model's documented API protocol. + + Args: + new_block: Block containing the request and generation settings. + conversation: Conversation history used to build the request. + system_message: Optional system instruction for the completion. + stop_block_index: Optional exclusive conversation-history boundary. + **kwargs: Additional protocol-specific completion options. + + Returns: + The selected protocol together with its raw completion response. + """ + protocol = self._protocol_for_model(new_block.model.model_id) + if protocol == _Protocol.RESPONSES: + value = self._responses_engine.completion( + new_block, + conversation, + system_message, + stop_block_index, + **kwargs, + ) + elif protocol == _Protocol.ANTHROPIC_MESSAGES: + value = self._anthropic_engine.completion( + new_block, + conversation, + system_message, + stop_block_index, + **kwargs, + ) + elif protocol == _Protocol.GEMINI: + value = self._gemini_engine.completion( + new_block, + conversation, + system_message, + stop_block_index, + **kwargs, + ) + else: + value = super().completion( + new_block, + conversation, + system_message, + stop_block_index, + **kwargs, + ) + return _ProtocolResponse(protocol=protocol, value=value) + + def completion_response_with_stream( + self, stream: _ProtocolResponse + ) -> Iterator[str | tuple[str, Any]]: + """Yield streaming response chunks through the selected protocol adapter. + + Args: + stream: Raw completion response paired with its selected protocol. + + Yields: + Text chunks or typed chunks such as citation tuples. + """ + if stream.protocol == _Protocol.RESPONSES: + return self._responses_engine.completion_response_with_stream( + stream.value + ) + if stream.protocol == _Protocol.ANTHROPIC_MESSAGES: + return self._anthropic_engine.completion_response_with_stream( + stream.value + ) + if stream.protocol == _Protocol.GEMINI: + return self._gemini_engine.completion_response_with_stream( + stream.value + ) + return super().completion_response_with_stream(stream.value) + + def cancel_completion(self, response: _ProtocolResponse) -> None: + """Cancel the active stream for every OpenCode protocol. + + Args: + response: Protocol-wrapped response to cancel. + """ + if response.protocol == _Protocol.GEMINI: + # A running Google Gen AI generator raises ``ValueError`` when its + # ``close`` method is called from another thread. Closing the client + # interrupts the HTTP stream; removing the cached client lets the next + # request create a fresh connection pool. + try: + super().cancel_completion(response.value) + except ValueError: + # Expected when the generator runs on another thread; closing the + # client below is what actually stops the stream. + log.debug("Gemini stream close deferred to client shutdown") + client = self._gemini_engine.__dict__.pop("client", None) + if client is not None: + client.close() + return + super().cancel_completion(response.value) + + def completion_response_without_stream( + self, response: _ProtocolResponse, new_block: MessageBlock, **kwargs + ) -> MessageBlock: + """Process a complete response through the selected protocol adapter. + + Args: + response: Raw completion response paired with its selected protocol. + new_block: Block to populate with the provider response. + **kwargs: Additional protocol-specific response-processing options. + + Returns: + The supplied block populated with the completed response. + """ + if response.protocol == _Protocol.RESPONSES: + return self._responses_engine.completion_response_without_stream( + response.value, new_block, **kwargs + ) + if response.protocol == _Protocol.ANTHROPIC_MESSAGES: + return self._anthropic_engine.completion_response_without_stream( + response.value, new_block, **kwargs + ) + if response.protocol == _Protocol.GEMINI: + return self._gemini_engine.completion_response_without_stream( + response.value, new_block, **kwargs + ) + return super().completion_response_without_stream( + response.value, new_block, **kwargs + ) + + +class OpenCodeGoEngine(_OpenCodeEngine): + """OpenCode Go subscription engine.""" + + VISION_MODELS = _GO_VISION_MODELS + FIXED_SAMPLING_PARAMETERS = _GO_FIXED_KIMI_SAMPLING_PARAMETERS + MODEL_UNSUPPORTED_PARAMETERS = _GO_MODEL_UNSUPPORTED_PARAMETERS + + +class OpenCodeZenEngine(_OpenCodeEngine): + """OpenCode Zen pay-as-you-go engine.""" + + MODEL_UNSUPPORTED_PARAMETERS = _ZEN_MODEL_UNSUPPORTED_PARAMETERS + VISION_MODELS = frozenset( + { + "gemini-3.7-flash", + "gemini-3.6-flash", + "gemini-3.5-flash", + "gemini-3.5-flash-lite", + "gemini-3.1-pro", + "gemini-3-flash", + "kimi-k2.7-code", + "qwen3.7-plus", + } + ) + + @staticmethod + def _protocol_for_model(model_id: str) -> _Protocol: + if model_id.startswith("gemini-"): + return _Protocol.GEMINI + if model_id.startswith(("claude-", "qwen")): + return _Protocol.ANTHROPIC_MESSAGES + if model_id.startswith(("gpt-", "grok-", "muse-")): + return _Protocol.RESPONSES + return _Protocol.CHAT_COMPLETIONS diff --git a/pyproject.toml b/pyproject.toml index 8d30c720..a59c8202 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -119,6 +119,9 @@ convention = "google" [tool.pytest.ini_options] +filterwarnings = [ + "ignore:'_UnionGenericAlias' is deprecated and slated for removal in Python 3.17:DeprecationWarning:google.genai.types", +] markers = [ "slow: marks tests as slow (deselect with '-m \"not slow\"')", "integration: marks tests as integration tests (deselect with '-m \"not integration\"')", diff --git a/tests/conftest.py b/tests/conftest.py index 56258eee..b79b7753 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -4,6 +4,7 @@ import pytest from PIL import Image +from pydantic_settings import YamlConfigSettingsSource from upath import UPath from basilisk.config.config_helper import BasiliskBaseSettings @@ -161,7 +162,8 @@ def mock_display_error_msg(mocker): def mock_settings_sources(mocker): """Mock the settings_customise_sources method to prevent loading real config files. - This overrides the method to only use init_settings, avoiding any file loading. + This uses init settings plus an empty YAML source, avoiding real file loading + while satisfying Pydantic's YAML configuration validation. """ @classmethod @@ -173,8 +175,10 @@ def mock_settings_customise_sources( dotenv_settings, file_secret_settings, ): - # Only use init_settings, skip loading from files - return (init_settings,) + return ( + init_settings, + YamlConfigSettingsSource(settings_cls, yaml_file=[]), + ) mocker.patch.object( BasiliskBaseSettings, diff --git a/tests/presenters/test_conversation_presenter.py b/tests/presenters/test_conversation_presenter.py index 5c6619df..1666f70b 100644 --- a/tests/presenters/test_conversation_presenter.py +++ b/tests/presenters/test_conversation_presenter.py @@ -178,6 +178,23 @@ def test_stops_running_completion(self, presenter, mocker): presenter.cleanup() mock_stop.assert_called_once_with(skip_callbacks=True) + def test_stops_dead_completion_with_logical_ownership( + self, presenter, mocker + ): + """cleanup() clears queued callbacks even after a worker has died.""" + mocker.patch.object( + presenter.completion_handler, "is_running", return_value=False + ) + mock_stop = mocker.patch.object( + presenter.completion_handler, "stop_completion" + ) + mocker.patch("basilisk.presenters.conversation_presenter.stop_sound") + mocker.patch.object(presenter, "flush_draft") + + presenter.cleanup() + + mock_stop.assert_called_once_with(skip_callbacks=True) + def test_aborts_live_recording_thread(self, presenter, mocker): """cleanup() aborts a live recording thread.""" mock_thread = MagicMock() diff --git a/tests/provider_engine/test_opencode_go_engine.py b/tests/provider_engine/test_opencode_go_engine.py new file mode 100644 index 00000000..131c969d --- /dev/null +++ b/tests/provider_engine/test_opencode_go_engine.py @@ -0,0 +1,483 @@ +"""Tests for the OpenCode and OpenCode Zen provider engines.""" + +from unittest.mock import MagicMock + +import httpx +import pytest +from google.genai.types import ResourceScope +from pydantic import SecretStr + +from basilisk.provider import ProviderAPIType, get_provider +from basilisk.provider_engine.opencode_go_engine import ( + OpenCodeGoEngine, + OpenCodeZenEngine, + _Protocol, + _ProtocolResponse, +) + +_LIVE_GO_MODEL_PROTOCOLS = ( + ("minimax-m3", _Protocol.ANTHROPIC_MESSAGES), + ("minimax-m2.7", _Protocol.ANTHROPIC_MESSAGES), + ("minimax-m2.5", _Protocol.ANTHROPIC_MESSAGES), + ("kimi-k3", _Protocol.CHAT_COMPLETIONS), + ("kimi-k2.7-code", _Protocol.CHAT_COMPLETIONS), + ("kimi-k2.6", _Protocol.CHAT_COMPLETIONS), + ("kimi-k2.5", _Protocol.CHAT_COMPLETIONS), + ("glm-5.2", _Protocol.CHAT_COMPLETIONS), + ("glm-5.3", _Protocol.CHAT_COMPLETIONS), + ("ox-alpha-free", _Protocol.CHAT_COMPLETIONS), + ("glm-5.1", _Protocol.CHAT_COMPLETIONS), + ("glm-5", _Protocol.CHAT_COMPLETIONS), + ("deepseek-v4-pro", _Protocol.CHAT_COMPLETIONS), + ("deepseek-v4-flash", _Protocol.CHAT_COMPLETIONS), + ("deepseek-v4-flash-vision-exp", _Protocol.CHAT_COMPLETIONS), + ("qwen3.7-max", _Protocol.ANTHROPIC_MESSAGES), + ("qwen3.8-max", _Protocol.ANTHROPIC_MESSAGES), + ("qwen3.7-plus", _Protocol.ANTHROPIC_MESSAGES), + ("qwen3.6-plus", _Protocol.ANTHROPIC_MESSAGES), + ("qwen3.5-plus", _Protocol.CHAT_COMPLETIONS), + ("mimo-v2-pro", _Protocol.CHAT_COMPLETIONS), + ("mimo-v2-omni", _Protocol.CHAT_COMPLETIONS), + ("mimo-v2.5-pro", _Protocol.CHAT_COMPLETIONS), + ("mimo-v2.5", _Protocol.CHAT_COMPLETIONS), + ("hy3", _Protocol.CHAT_COMPLETIONS), + ("hy3-preview", _Protocol.CHAT_COMPLETIONS), + ("gpt-5.6-luna", _Protocol.RESPONSES), + ("grok-4.5", _Protocol.RESPONSES), + ("muse-spark-1.2-contributor", _Protocol.RESPONSES), +) +_LIVE_GO_MODEL_VISION = ( + ("minimax-m3", False), + ("minimax-m2.7", False), + ("minimax-m2.5", False), + ("kimi-k3", True), + ("kimi-k2.7-code", True), + ("kimi-k2.6", True), + ("kimi-k2.5", True), + ("glm-5.2", False), + ("glm-5.3", False), + ("ox-alpha-free", True), + ("glm-5.1", False), + ("glm-5", False), + ("deepseek-v4-pro", False), + ("deepseek-v4-flash", False), + ("deepseek-v4-flash-vision-exp", True), + ("qwen3.7-max", False), + ("qwen3.8-max", True), + ("qwen3.7-plus", True), + ("qwen3.6-plus", True), + ("qwen3.5-plus", True), + ("mimo-v2-pro", True), + ("mimo-v2-omni", True), + ("mimo-v2.5-pro", True), + ("mimo-v2.5", True), + ("hy3", False), + ("hy3-preview", False), + ("gpt-5.6-luna", True), + ("grok-4.5", True), + ("muse-spark-1.2-contributor", True), +) + + +def _make_engine(engine_cls, provider_id): + account = MagicMock() + account.id = "account-test" + account.provider = get_provider(id=provider_id) + account.api_key = SecretStr("sk-test") + account.custom_base_url = None + return engine_cls(account) + + +@pytest.mark.parametrize( + ("provider_id", "name", "base_url", "engine_cls"), + [ + ( + "opencodego", + "OpenCode Go", + "https://opencode.ai/zen/go/v1", + OpenCodeGoEngine, + ), + ( + "opencodezen", + "OpenCode Zen", + "https://opencode.ai/zen/v1", + OpenCodeZenEngine, + ), + ], +) +def test_provider_registration(provider_id, name, base_url, engine_cls): + """Both account options expose their intended configuration.""" + provider = get_provider(id=provider_id) + + assert provider.name == name + assert provider.base_url == base_url + assert provider.api_type == ProviderAPIType.OPENAI + assert provider.env_var_name_api_key == "OPENCODE_API_KEY" + assert provider.engine_cls is engine_cls + + +@pytest.mark.parametrize( + ("engine_cls", "provider_id", "models_url", "expected_kimi_extra_info"), + [ + ( + OpenCodeGoEngine, + "opencodego", + "https://opencode.ai/zen/go/v1/models", + { + "owned_by": "opencode", + "unsupported_parameters": ["temperature", "top_p"], + }, + ), + ( + OpenCodeZenEngine, + "opencodezen", + "https://opencode.ai/zen/v1/models", + { + "owned_by": "opencode", + "unsupported_parameters": ["temperature", "top_p"], + }, + ), + ], +) +def test_model_discovery( + httpx_mock, engine_cls, provider_id, models_url, expected_kimi_extra_info +): + """Each product discovers models from its own endpoint.""" + httpx_mock.add_response( + url=models_url, + json={ + "data": [ + {"id": "qwen3.7-plus", "created": 123, "owned_by": "opencode"}, + {"id": "kimi-k3", "created": 124, "owned_by": "opencode"}, + {"id": "text-model", "created": "invalid"}, + {"created": 456}, + "invalid", + ] + }, + ) + engine = _make_engine(engine_cls, provider_id) + + models = engine._load_models() + + assert [model.id for model in models] == [ + "qwen3.7-plus", + "kimi-k3", + "text-model", + ] + assert models[0].vision is True + assert models[0].created == 123 + assert models[0].extra_info == {"owned_by": "opencode"} + assert models[1].extra_info == expected_kimi_extra_info + assert models[2].created == 0 + request = httpx_mock.get_requests()[0] + assert request.headers["Authorization"] == "Bearer sk-test" + + +def test_model_discovery_rejects_non_list(httpx_mock): + """Malformed model payloads report a clear error.""" + httpx_mock.add_response( + url="https://opencode.ai/zen/v1/models", json={"data": {}} + ) + engine = _make_engine(OpenCodeZenEngine, "opencodezen") + + with pytest.raises(ValueError, match="invalid model list"): + engine._load_models() + + +def test_model_discovery_http_error(httpx_mock): + """Model-list HTTP failures preserve status and response details.""" + httpx_mock.add_response( + url="https://opencode.ai/zen/go/v1/models", + status_code=401, + text="invalid key", + ) + engine = _make_engine(OpenCodeGoEngine, "opencodego") + + with pytest.raises(httpx.HTTPStatusError, match="status=401.*invalid key"): + engine._load_models() + + +@pytest.mark.parametrize(("model_id", "protocol"), _LIVE_GO_MODEL_PROTOCOLS) +def test_live_go_models_use_audited_protocol(model_id, protocol): + """Every live Go model uses its audited transport.""" + engine = _make_engine(OpenCodeGoEngine, "opencodego") + + assert engine._protocol_for_model(model_id) == protocol + + +def test_unknown_go_model_uses_chat_completions(): + """Future Go models retain the documented conservative fallback.""" + engine = _make_engine(OpenCodeGoEngine, "opencodego") + + assert ( + engine._protocol_for_model("future-model") == _Protocol.CHAT_COMPLETIONS + ) + + +@pytest.mark.parametrize( + ("model_id", "protocol"), + [ + ("gemini-3.7-flash", _Protocol.GEMINI), + ("claude-sonnet-5", _Protocol.ANTHROPIC_MESSAGES), + ("gpt-5.6-sol", _Protocol.RESPONSES), + ("deepseek-v4-pro", _Protocol.CHAT_COMPLETIONS), + ], +) +def test_zen_protocol_selection_is_unchanged(model_id, protocol): + """The Go audit does not alter Zen routing.""" + engine = _make_engine(OpenCodeZenEngine, "opencodezen") + + assert engine._protocol_for_model(model_id) == protocol + + +@pytest.mark.parametrize(("model_id", "vision"), _LIVE_GO_MODEL_VISION) +def test_live_go_models_use_audited_vision_metadata(model_id, vision): + """Every audited live Go model has its catalogued attachment support.""" + engine = _make_engine(OpenCodeGoEngine, "opencodego") + + assert (model_id in engine.VISION_MODELS) is vision + + +@pytest.mark.parametrize( + ("protocol", "adapter_name"), + [ + (_Protocol.RESPONSES, "_responses_engine"), + (_Protocol.ANTHROPIC_MESSAGES, "_anthropic_engine"), + (_Protocol.GEMINI, "_gemini_engine"), + ], +) +def test_completion_routes_to_adapter(mocker, protocol, adapter_name): + """Completion creation delegates to the selected protocol adapter.""" + engine = _make_engine(OpenCodeZenEngine, "opencodezen") + adapter = MagicMock() + adapter.completion.return_value = "raw-response" + mocker.patch.object(engine, adapter_name, adapter) + mocker.patch.object(engine, "_protocol_for_model", return_value=protocol) + new_block = MagicMock() + new_block.model.model_id = "model-test" + conversation = MagicMock() + + response = engine.completion( + new_block, conversation, None, stop_block_index=2, extra="value" + ) + + assert response == _ProtocolResponse(protocol, "raw-response") + adapter.completion.assert_called_once_with( + new_block, conversation, None, 2, extra="value" + ) + + +@pytest.mark.parametrize( + ("protocol", "adapter_name"), + [ + (_Protocol.RESPONSES, "_responses_engine"), + (_Protocol.ANTHROPIC_MESSAGES, "_anthropic_engine"), + (_Protocol.GEMINI, "_gemini_engine"), + ], +) +def test_response_processing_routes_to_adapter(mocker, protocol, adapter_name): + """Streaming and complete responses use the selected adapter.""" + engine = _make_engine(OpenCodeZenEngine, "opencodezen") + adapter = MagicMock() + adapter.completion_response_with_stream.return_value = iter(["chunk"]) + completed_block = MagicMock() + adapter.completion_response_without_stream.return_value = completed_block + mocker.patch.object(engine, adapter_name, adapter) + response = _ProtocolResponse(protocol, "raw-response") + new_block = MagicMock() + + assert list(engine.completion_response_with_stream(response)) == ["chunk"] + assert ( + engine.completion_response_without_stream(response, new_block) + is completed_block + ) + adapter.completion_response_with_stream.assert_called_once_with( + "raw-response" + ) + adapter.completion_response_without_stream.assert_called_once_with( + "raw-response", new_block + ) + + +@pytest.mark.parametrize( + "model_id", ["kimi-k2.5", "kimi-k2.6", "kimi-k2.7-code", "kimi-k3"] +) +def test_go_kimi_models_use_fixed_sampling_parameters( + httpx_mock, mocker, model_id +): + """Go Kimi chat requests override stored sampling values.""" + engine = _make_engine(OpenCodeGoEngine, "opencodego") + httpx_mock.add_response( + url="https://opencode.ai/zen/go/v1/models", + json={"data": [{"id": model_id}]}, + ) + model = engine._load_models()[0] + assert model.extra_info["unsupported_parameters"] == [ + "temperature", + "top_p", + ] + client = MagicMock() + mocker.patch.object(engine, "get_model", return_value=model) + mocker.patch.object(engine, "get_messages", return_value=[]) + mocker.patch.object(engine, "client", client) + new_block = MagicMock() + new_block.model.model_id = model_id + new_block.stream = False + new_block.temperature = 0.7 + new_block.top_p = 0.8 + new_block.max_tokens = 100 + + engine.completion(new_block, MagicMock(), None) + + params = client.chat.completions.create.call_args.kwargs + assert params["temperature"] == 1.0 + assert params["top_p"] == 0.95 + + +def test_go_luna_responses_request_omits_sampling_parameters( + httpx_mock, mocker +): + """GPT-5.6 Luna omits unsupported Responses sampling parameters.""" + engine = _make_engine(OpenCodeGoEngine, "opencodego") + httpx_mock.add_response( + url="https://opencode.ai/zen/go/v1/models", + json={"data": [{"id": "gpt-5.6-luna"}]}, + ) + model = engine._load_models()[0] + assert model.extra_info["unsupported_parameters"] == [ + "temperature", + "top_p", + ] + adapter = engine._responses_engine + client = MagicMock() + mocker.patch.object(adapter, "get_model", return_value=model) + mocker.patch.object(adapter, "get_messages", return_value=[]) + mocker.patch.object(adapter, "client", client) + new_block = MagicMock() + new_block.model.model_id = model.id + new_block.stream = False + new_block.temperature = 0.7 + new_block.top_p = 0.8 + new_block.max_tokens = 100 + + engine.completion(new_block, MagicMock(), None) + + params = client.responses.create.call_args.kwargs + assert "temperature" not in params + assert "top_p" not in params + + +def test_ordinary_go_chat_model_keeps_user_sampling_parameters( + httpx_mock, mocker +): + """Go models without a policy retain the user's selected sampling values.""" + engine = _make_engine(OpenCodeGoEngine, "opencodego") + httpx_mock.add_response( + url="https://opencode.ai/zen/go/v1/models", + json={"data": [{"id": "text-model"}]}, + ) + model = engine._load_models()[0] + client = MagicMock() + mocker.patch.object(engine, "get_model", return_value=model) + mocker.patch.object(engine, "get_messages", return_value=[]) + mocker.patch.object(engine, "client", client) + new_block = MagicMock() + new_block.model.model_id = model.id + new_block.stream = False + new_block.temperature = 0.7 + new_block.top_p = 0.8 + new_block.max_tokens = 100 + + engine.completion(new_block, MagicMock(), None) + + params = client.chat.completions.create.call_args.kwargs + assert params["temperature"] == 0.7 + assert params["top_p"] == 0.8 + + +def test_go_fixed_sampling_policy_does_not_apply_to_zen(httpx_mock, mocker): + """Zen Kimi K3 keeps its independent unsupported-parameter policy.""" + engine = _make_engine(OpenCodeZenEngine, "opencodezen") + httpx_mock.add_response( + url="https://opencode.ai/zen/v1/models", + json={"data": [{"id": "kimi-k3"}]}, + ) + model = engine._load_models()[0] + assert model.extra_info["unsupported_parameters"] == [ + "temperature", + "top_p", + ] + client = MagicMock() + mocker.patch.object(engine, "get_model", return_value=model) + mocker.patch.object(engine, "get_messages", return_value=[]) + mocker.patch.object(engine, "client", client) + new_block = MagicMock() + new_block.model.model_id = model.id + new_block.stream = False + new_block.temperature = 0.7 + new_block.top_p = 0.8 + new_block.max_tokens = 100 + + engine.completion(new_block, MagicMock(), None) + + params = client.chat.completions.create.call_args.kwargs + assert "temperature" not in params + assert "top_p" not in params + + +def test_protocol_client_base_urls(): + """All protocol clients target the OpenCode Zen gateway.""" + engine = _make_engine(OpenCodeZenEngine, "opencodezen") + + assert str(engine.client.base_url) == "https://opencode.ai/zen/v1/" + assert str(engine._responses_engine.client.base_url) == ( + "https://opencode.ai/zen/v1/" + ) + assert str(engine._anthropic_engine.client.base_url) == ( + "https://opencode.ai/zen/" + ) + assert ( + engine._gemini_engine.client._api_client._http_options.base_url + == "https://opencode.ai/zen/v1" + ) + assert ( + engine._gemini_engine.client._api_client._http_options.base_url_resource_scope + == ResourceScope.COLLECTION + ) + + +@pytest.mark.parametrize( + ("engine_cls", "provider_id"), + [(OpenCodeGoEngine, "opencodego"), (OpenCodeZenEngine, "opencodezen")], +) +@pytest.mark.parametrize( + "protocol", + [ + _Protocol.CHAT_COMPLETIONS, + _Protocol.RESPONSES, + _Protocol.ANTHROPIC_MESSAGES, + ], +) +def test_cancel_closes_go_and_zen_streams(engine_cls, provider_id, protocol): + """Go and Zen close OpenAI-compatible and Anthropic SDK streams.""" + engine = _make_engine(engine_cls, provider_id) + stream = MagicMock() + + engine.cancel_completion(_ProtocolResponse(protocol, stream)) + + stream.close.assert_called_once_with() + + +def test_cancel_zen_gemini_stream_closes_and_recreates_client(): + """Zen cancels Gemini's generator by closing its active HTTP client.""" + engine = _make_engine(OpenCodeZenEngine, "opencodezen") + stream = MagicMock() + stream.close.side_effect = ValueError("generator already executing") + client = MagicMock() + engine._gemini_engine.__dict__["client"] = client + + engine.cancel_completion(_ProtocolResponse(_Protocol.GEMINI, stream)) + + stream.close.assert_called_once_with() + client.close.assert_called_once_with() + assert "client" not in engine._gemini_engine.__dict__ diff --git a/tests/test_completion_handler.py b/tests/test_completion_handler.py new file mode 100644 index 00000000..f9b6b8d6 --- /dev/null +++ b/tests/test_completion_handler.py @@ -0,0 +1,973 @@ +"""Tests for completion lifecycle and cancellation.""" + +from __future__ import annotations + +import threading +from unittest.mock import MagicMock + +import pytest + +from basilisk.completion_handler import CompletionHandler + + +class _BlockingStream: + """A stream that stays blocked until the provider response is closed.""" + + def __init__(self): + self.read_started = threading.Event() + self.closed = threading.Event() + + def __iter__(self): + return self + + def __next__(self): + self.read_started.set() + if not self.closed.wait(timeout=2): + raise TimeoutError("test stream was not cancelled") + raise OSError("stream closed") + + def close(self): + self.closed.set() + + +class _PartialThenBlockingStream: + """A stream that yields partial text before waiting to be cancelled.""" + + def __init__(self, partial_text): + self.partial_text = partial_text + self.read_started = threading.Event() + self.closed = threading.Event() + self._yielded_partial = False + + def __iter__(self): + return self + + def __next__(self): + if not self._yielded_partial: + self._yielded_partial = True + return self.partial_text + self.read_started.set() + if not self.closed.wait(timeout=2): + raise TimeoutError("test stream was not cancelled") + raise OSError("stream closed") + + def close(self): + self.closed.set() + + +def test_stop_closes_blocked_provider_stream( + mocker, empty_conversation, message_block +): + """Stop interrupts a stream without waiting for another server chunk.""" + stream = _BlockingStream() + engine = MagicMock() + engine.completion.return_value = stream + engine.completion_response_with_stream.return_value = stream + engine.cancel_completion.side_effect = lambda response: response.close() + completion_end = MagicMock() + on_error = MagicMock() + mocker.patch( + "basilisk.completion_handler.wx.CallAfter", + side_effect=lambda callback, *args: callback(*args), + ) + mocker.patch("basilisk.completion_handler.play_sound") + stop_sound = mocker.patch("basilisk.completion_handler.stop_sound") + handler = CompletionHandler( + on_completion_end=completion_end, on_error=on_error + ) + + handler.start_completion( + engine=engine, + system_message=None, + conversation=empty_conversation, + new_block=message_block, + stream=True, + ) + assert stream.read_started.wait(timeout=1) + assert handler.is_running() + assert handler._active_response is stream + + handler.stop_completion() + + engine.cancel_completion.assert_called_once_with(stream) + assert stream.closed.wait(timeout=1) + assert not handler.is_running() + stop_sound.assert_called_once_with() + completion_end.assert_called_once_with(False) + on_error.assert_not_called() + + +def test_completion_start_callback_stops_published_blocking_worker( + mocker, empty_conversation, message_block +): + """An immediate start callback cancels its published worker cleanly.""" + stream = _BlockingStream() + engine = MagicMock() + engine.completion.return_value = stream + engine.completion_response_with_stream.return_value = stream + engine.cancel_completion.side_effect = lambda response: response.close() + completion_end = MagicMock() + on_error = MagicMock() + on_stream_start = MagicMock() + on_stream_chunk = MagicMock() + on_stream_finish = MagicMock() + published_tasks = [] + mocker.patch( + "basilisk.completion_handler.wx.CallAfter", + side_effect=lambda callback, *args: callback(*args), + ) + mocker.patch("basilisk.completion_handler.play_sound") + mocker.patch("basilisk.completion_handler.stop_sound") + + def stop_on_completion_start(): + published_tasks.append(handler.task) + handler.stop_completion() + + handler = CompletionHandler( + on_completion_start=stop_on_completion_start, + on_completion_end=completion_end, + on_error=on_error, + on_stream_start=on_stream_start, + on_stream_chunk=on_stream_chunk, + on_stream_finish=on_stream_finish, + ) + + handler.start_completion( + engine=engine, + system_message=None, + conversation=empty_conversation, + new_block=message_block, + stream=True, + ) + + assert len(published_tasks) == 1 + published_task = published_tasks[0] + assert published_task is not None + published_task.join(timeout=1) + assert not published_task.is_alive() + assert not handler.is_running() + assert handler.task is None + assert handler._active_request is None + assert handler._active_engine is None + assert handler._active_response is None + assert not handler._stream_buffers + engine.completion.assert_not_called() + engine.cancel_completion.assert_not_called() + assert not stream.read_started.is_set() + assert not stream.closed.is_set() + completion_end.assert_called_once_with(False) + on_error.assert_not_called() + on_stream_start.assert_not_called() + on_stream_chunk.assert_not_called() + on_stream_finish.assert_not_called() + + +def test_completion_start_scheduled_stop_closes_published_blocking_response( + mocker, empty_conversation, message_block +): + """A start-lifecycle stop closes a response published after the start gate.""" + queued_callbacks = [] + stream = _BlockingStream() + engine = MagicMock() + engine.completion.return_value = stream + engine.completion_response_with_stream.return_value = stream + engine.cancel_completion.side_effect = lambda response: response.close() + completion_end = MagicMock() + on_error = MagicMock() + on_stream_start = MagicMock() + on_stream_chunk = MagicMock() + on_stream_finish = MagicMock() + published_tasks = [] + helper_tasks = [] + mocker.patch( + "basilisk.completion_handler.wx.CallAfter", + side_effect=lambda callback, *args: queued_callbacks.append( + (callback, args) + ), + ) + mocker.patch("basilisk.completion_handler.play_sound") + mocker.patch("basilisk.completion_handler.stop_sound") + + def stop_after_stream_starts(): + assert stream.read_started.wait(timeout=1) + handler.stop_completion() + + def schedule_stop_on_completion_start(): + published_tasks.append(handler.task) + helper = threading.Thread(target=stop_after_stream_starts) + helper_tasks.append(helper) + helper.start() + + handler = CompletionHandler( + on_completion_start=schedule_stop_on_completion_start, + on_completion_end=completion_end, + on_error=on_error, + on_stream_start=on_stream_start, + on_stream_chunk=on_stream_chunk, + on_stream_finish=on_stream_finish, + ) + + handler.start_completion( + engine=engine, + system_message=None, + conversation=empty_conversation, + new_block=message_block, + stream=True, + ) + + assert len(published_tasks) == 1 + assert len(helper_tasks) == 1 + published_task = published_tasks[0] + helper_task = helper_tasks[0] + assert published_task is not None + helper_task.join(timeout=1) + published_task.join(timeout=1) + assert not helper_task.is_alive() + assert not published_task.is_alive() + assert stream.closed.is_set() + engine.cancel_completion.assert_called_once_with(stream) + assert handler.task is None + assert handler._active_request is None + assert handler._active_engine is None + assert handler._active_response is None + assert not handler._stream_buffers + + for callback, args in queued_callbacks: + callback(*args) + + completion_end.assert_called_once_with(False) + on_error.assert_not_called() + on_stream_start.assert_not_called() + on_stream_chunk.assert_not_called() + on_stream_finish.assert_not_called() + + +def test_stop_between_start_gate_and_startup_commit_aborts_provider_start( + mocker, empty_conversation, message_block +): + """Stop wins deterministically after the gate but before startup commits.""" + startup_waiting = threading.Event() + release_startup = threading.Event() + published_tasks = [] + engine = MagicMock() + completion_end = MagicMock() + mocker.patch( + "basilisk.completion_handler.wx.CallAfter", + side_effect=lambda callback, *args: callback(*args), + ) + play_sound = mocker.patch("basilisk.completion_handler.play_sound") + mocker.patch("basilisk.completion_handler.stop_sound") + handler = CompletionHandler( + on_completion_start=lambda: published_tasks.append(handler.task), + on_completion_end=completion_end, + ) + original_commit_startup = handler._commit_startup + + def wait_before_startup_commit(request): + startup_waiting.set() + assert release_startup.wait(timeout=1) + return original_commit_startup(request) + + mocker.patch.object( + handler, "_commit_startup", side_effect=wait_before_startup_commit + ) + handler.start_completion( + engine=engine, + system_message=None, + conversation=empty_conversation, + new_block=message_block, + ) + assert startup_waiting.wait(timeout=1) + + handler.stop_completion() + release_startup.set() + + assert len(published_tasks) == 1 + published_task = published_tasks[0] + assert published_task is not None + published_task.join(timeout=1) + assert not published_task.is_alive() + play_sound.assert_not_called() + engine.completion.assert_not_called() + assert handler.task is None + assert handler._active_request is None + assert handler._active_engine is None + assert handler._active_response is None + assert not handler._stream_buffers + completion_end.assert_called_once_with(False) + + +def test_fast_completion_notifies_start_before_end( + mocker, empty_conversation, message_block +): + """A fast worker cannot finish before the start callback runs.""" + callback_order = [] + completion_finished = threading.Event() + engine = MagicMock() + engine.completion.return_value = object() + engine.completion_response_without_stream.side_effect = lambda **kwargs: ( + kwargs["new_block"] + ) + mocker.patch( + "basilisk.completion_handler.wx.CallAfter", + side_effect=lambda callback, *args: callback(*args), + ) + mocker.patch("basilisk.completion_handler.play_sound") + mocker.patch("basilisk.completion_handler.stop_sound") + + def on_completion_end(success): + callback_order.append(("end", success)) + completion_finished.set() + + handler = CompletionHandler( + on_completion_start=lambda: callback_order.append("start"), + on_completion_end=on_completion_end, + ) + + handler.start_completion( + engine=engine, + system_message=None, + conversation=empty_conversation, + new_block=message_block, + ) + + assert completion_finished.wait(timeout=1) + assert callback_order == ["start", ("end", True)] + + +@pytest.mark.parametrize( + ("failure_site", "failure_message"), + [("provider", "provider failed"), ("sound", "sound failed")], +) +def test_start_callback_precedes_fast_start_failure_callbacks( + mocker, empty_conversation, message_block, failure_site, failure_message +): + """Provider-start errors cannot escape while start notification is running.""" + callback_order = [] + start_entered = threading.Event() + release_start = threading.Event() + completion_finished = threading.Event() + engine = MagicMock() + play_sound = mocker.patch("basilisk.completion_handler.play_sound") + if failure_site == "provider": + engine.completion.side_effect = RuntimeError(failure_message) + else: + play_sound.side_effect = RuntimeError(failure_message) + + def on_completion_start(): + callback_order.append("start") + start_entered.set() + assert release_start.wait(timeout=1) + callback_order.append("start returned") + + def on_completion_end(success): + callback_order.append(("end", success)) + completion_finished.set() + + handler = CompletionHandler( + on_completion_start=on_completion_start, + on_completion_end=on_completion_end, + on_error=lambda error: callback_order.append(("error", error)), + ) + mocker.patch( + "basilisk.completion_handler.wx.CallAfter", + side_effect=lambda callback, *args: callback(*args), + ) + mocker.patch("basilisk.completion_handler.stop_sound") + + start_thread = threading.Thread( + target=handler.start_completion, + kwargs={ + "engine": engine, + "system_message": None, + "conversation": empty_conversation, + "new_block": message_block, + }, + ) + start_thread.start() + assert start_entered.wait(timeout=1) + assert callback_order == ["start"] + + release_start.set() + start_thread.join(timeout=1) + assert not start_thread.is_alive() + assert completion_finished.wait(timeout=1) + assert callback_order == [ + "start", + "start returned", + ("error", failure_message), + ("end", False), + ] + + +def test_stop_skip_callbacks_suppresses_queued_start_error( + mocker, empty_conversation, message_block +): + """Destroy-time stop suppresses an error queued by an already-dead worker.""" + queued_callbacks = [] + engine = MagicMock() + engine.completion.side_effect = RuntimeError("provider failed") + completion_end = MagicMock() + on_error = MagicMock() + mocker.patch( + "basilisk.completion_handler.wx.CallAfter", + side_effect=lambda callback, *args: queued_callbacks.append( + (callback, args) + ), + ) + mocker.patch("basilisk.completion_handler.play_sound") + stop_sound = mocker.patch("basilisk.completion_handler.stop_sound") + handler = CompletionHandler( + on_completion_end=completion_end, on_error=on_error + ) + + handler.start_completion( + engine=engine, + system_message=None, + conversation=empty_conversation, + new_block=message_block, + ) + task = handler.task + assert task is not None + task.join(timeout=1) + assert not task.is_alive() + assert len(queued_callbacks) == 1 + + handler.stop_completion(skip_callbacks=True) + callback, args = queued_callbacks.pop() + callback(*args) + + assert handler.task is None + assert handler._active_request is None + assert handler._active_engine is None + assert handler._active_response is None + assert handler._latest_request is None + assert not handler._stream_buffers + on_error.assert_not_called() + completion_end.assert_not_called() + stop_sound.assert_called_once_with() + + +def test_thread_start_failure_clears_completion_state_for_retry( + mocker, empty_conversation, message_block +): + """A thread start failure rolls back publication before a later retry.""" + start_calls = 0 + completion_finished = threading.Event() + completion_start = MagicMock() + completion_end = MagicMock() + on_error = MagicMock() + engine = MagicMock() + engine.completion.return_value = object() + engine.completion_response_without_stream.side_effect = lambda **kwargs: ( + kwargs["new_block"] + ) + original_thread_start = threading.Thread.start + + def fail_first_thread_start(task): + nonlocal start_calls + start_calls += 1 + if start_calls == 1: + raise RuntimeError("thread failed to start") + return original_thread_start(task) + + mocker.patch.object( + threading.Thread, + "start", + autospec=True, + side_effect=fail_first_thread_start, + ) + mocker.patch( + "basilisk.completion_handler.wx.CallAfter", + side_effect=lambda callback, *args: callback(*args), + ) + mocker.patch("basilisk.completion_handler.play_sound") + mocker.patch("basilisk.completion_handler.stop_sound") + + def on_completion_end(success): + completion_end(success) + completion_finished.set() + + handler = CompletionHandler( + on_completion_start=completion_start, + on_completion_end=on_completion_end, + on_error=on_error, + ) + + with pytest.raises(RuntimeError, match="thread failed to start"): + handler.start_completion( + engine=engine, + system_message=None, + conversation=empty_conversation, + new_block=message_block, + ) + + assert handler.task is None + assert handler._active_request is None + assert handler._latest_request is None + assert handler._active_engine is None + assert handler._active_response is None + assert handler._start_notified is None + assert not handler._stream_buffers + completion_start.assert_not_called() + completion_end.assert_not_called() + on_error.assert_not_called() + + handler.start_completion( + engine=engine, + system_message=None, + conversation=empty_conversation, + new_block=message_block, + ) + assert completion_finished.wait(timeout=1) + assert start_calls == 2 + completion_start.assert_called_once_with() + completion_end.assert_called_once_with(True) + on_error.assert_not_called() + + +def test_completion_start_failure_cancels_and_cleans_published_worker( + mocker, empty_conversation, message_block +): + """A failing start callback suppresses the worker and re-raises its error.""" + engine = MagicMock() + completion_end = MagicMock() + on_error = MagicMock() + on_stream_start = MagicMock() + on_stream_chunk = MagicMock() + on_stream_finish = MagicMock() + callback_error = RuntimeError("start callback failed") + mocker.patch( + "basilisk.completion_handler.wx.CallAfter", + side_effect=lambda callback, *args: callback(*args), + ) + play_sound = mocker.patch("basilisk.completion_handler.play_sound") + mocker.patch("basilisk.completion_handler.stop_sound") + + def fail_completion_start(): + raise callback_error + + handler = CompletionHandler( + on_completion_start=fail_completion_start, + on_completion_end=completion_end, + on_error=on_error, + on_stream_start=on_stream_start, + on_stream_chunk=on_stream_chunk, + on_stream_finish=on_stream_finish, + ) + + with pytest.raises(RuntimeError, match="start callback failed"): + handler.start_completion( + engine=engine, + system_message=None, + conversation=empty_conversation, + new_block=message_block, + ) + + assert handler.task is None + assert handler._active_request is None + assert handler._active_engine is None + assert handler._active_response is None + assert handler._latest_request is None + assert not handler._stream_buffers + engine.completion.assert_not_called() + play_sound.assert_not_called() + completion_end.assert_not_called() + on_error.assert_not_called() + on_stream_start.assert_not_called() + on_stream_chunk.assert_not_called() + on_stream_finish.assert_not_called() + + +def test_stale_worker_cleanup_preserves_newer_active_response( + mocker, empty_conversation, message_block +): + """An old worker's finally block cannot clear a newer response.""" + processing_started = threading.Event() + allow_old_worker_to_finish = threading.Event() + old_response = object() + new_response = object() + old_engine = MagicMock() + old_engine.completion.return_value = old_response + + def process_old_response(**kwargs): + processing_started.set() + assert allow_old_worker_to_finish.wait(timeout=1) + return kwargs["new_block"] + + old_engine.completion_response_without_stream.side_effect = ( + process_old_response + ) + mocker.patch("basilisk.completion_handler.play_sound") + mocker.patch("basilisk.completion_handler.wx.CallAfter") + handler = CompletionHandler() + old_request = object() + old_task = threading.Thread( + target=handler._handle_completion, + kwargs={ + "request": old_request, + "engine": old_engine, + "system_message": None, + "conversation": empty_conversation, + "new_block": message_block, + "stream": False, + }, + ) + with handler._completion_lock: + handler._active_request = old_request + handler.task = old_task + old_task.start() + assert processing_started.wait(timeout=1) + + new_request = object() + new_engine = object() + new_task = object() + with handler._completion_lock: + handler._active_request = new_request + handler._active_engine = new_engine + handler._active_response = new_response + handler.task = new_task + + allow_old_worker_to_finish.set() + old_task.join(timeout=1) + assert not old_task.is_alive() + assert handler._active_request is new_request + assert handler._active_engine is new_engine + assert handler._active_response is new_response + assert handler.task is new_task + + +def test_stale_success_callback_does_not_finish_newer_completion( + mocker, empty_conversation, message_block +): + """A queued success callback cannot clear a replacement worker.""" + queued_callbacks = [] + mocker.patch( + "basilisk.completion_handler.wx.CallAfter", + side_effect=lambda callback, *args: queued_callbacks.append( + (callback, args) + ), + ) + mocker.patch("basilisk.completion_handler.play_sound") + mocker.patch("basilisk.completion_handler.stop_sound") + completion_end = MagicMock() + handler = CompletionHandler(on_completion_end=completion_end) + first_engine = MagicMock() + first_engine.completion.return_value = object() + first_engine.completion_response_without_stream.side_effect = ( + lambda **kwargs: kwargs["new_block"] + ) + + handler.start_completion( + engine=first_engine, + system_message=None, + conversation=empty_conversation, + new_block=message_block, + ) + first_task = handler.task + assert first_task is not None + first_task.join(timeout=1) + assert not first_task.is_alive() + assert len(queued_callbacks) == 1 + + second_started = threading.Event() + allow_second_completion = threading.Event() + second_engine = MagicMock() + + def start_second_completion(**kwargs): + second_started.set() + assert allow_second_completion.wait(timeout=1) + return object() + + second_engine.completion.side_effect = start_second_completion + second_engine.completion_response_without_stream.side_effect = ( + lambda **kwargs: kwargs["new_block"] + ) + handler.start_completion( + engine=second_engine, + system_message=None, + conversation=empty_conversation, + new_block=message_block, + ) + assert second_started.wait(timeout=1) + second_task = handler.task + assert second_task is not None + + callback, args = queued_callbacks.pop(0) + callback(*args) + assert handler.task is second_task + completion_end.assert_not_called() + + allow_second_completion.set() + second_task.join(timeout=1) + callback, args = queued_callbacks.pop(0) + callback(*args) + completion_end.assert_called_once_with(True) + assert handler.task is None + + +def test_stale_error_callback_does_not_finish_newer_completion( + mocker, empty_conversation, message_block +): + """A queued error callback cannot clear or fail a replacement worker.""" + queued_callbacks = [] + mocker.patch( + "basilisk.completion_handler.wx.CallAfter", + side_effect=lambda callback, *args: queued_callbacks.append( + (callback, args) + ), + ) + mocker.patch("basilisk.completion_handler.play_sound") + mocker.patch("basilisk.completion_handler.stop_sound") + completion_end = MagicMock() + on_error = MagicMock() + handler = CompletionHandler( + on_completion_end=completion_end, on_error=on_error + ) + first_engine = MagicMock() + first_engine.completion.side_effect = RuntimeError("first failure") + + handler.start_completion( + engine=first_engine, + system_message=None, + conversation=empty_conversation, + new_block=message_block, + ) + first_task = handler.task + assert first_task is not None + first_task.join(timeout=1) + assert not first_task.is_alive() + assert len(queued_callbacks) == 1 + + second_started = threading.Event() + allow_second_completion = threading.Event() + second_engine = MagicMock() + + def start_second_completion(**kwargs): + second_started.set() + assert allow_second_completion.wait(timeout=1) + return object() + + second_engine.completion.side_effect = start_second_completion + second_engine.completion_response_without_stream.side_effect = ( + lambda **kwargs: kwargs["new_block"] + ) + handler.start_completion( + engine=second_engine, + system_message=None, + conversation=empty_conversation, + new_block=message_block, + ) + assert second_started.wait(timeout=1) + second_task = handler.task + assert second_task is not None + + callback, args = queued_callbacks.pop(0) + callback(*args) + assert handler.task is second_task + on_error.assert_not_called() + completion_end.assert_not_called() + + allow_second_completion.set() + second_task.join(timeout=1) + callback, args = queued_callbacks.pop(0) + callback(*args) + completion_end.assert_called_once_with(True) + assert handler.task is None + + +def test_stopped_stream_partial_buffer_does_not_leak_to_restart( + mocker, empty_conversation, message_block +): + """Stopping a stream discards its partial text before the next request.""" + old_stream = _PartialThenBlockingStream("old partial") + old_engine = MagicMock() + old_engine.completion.return_value = old_stream + old_engine.completion_response_with_stream.return_value = old_stream + old_engine.cancel_completion.side_effect = lambda response: response.close() + new_engine = MagicMock() + new_engine.completion.return_value = iter(["new response."]) + new_engine.completion_response_with_stream.side_effect = lambda response: ( + response + ) + stream_chunks = MagicMock() + completion_results = [] + new_completion_finished = threading.Event() + + def on_completion_end(success): + completion_results.append(success) + if success: + new_completion_finished.set() + + mocker.patch( + "basilisk.completion_handler.wx.CallAfter", + side_effect=lambda callback, *args: callback(*args), + ) + mocker.patch("basilisk.completion_handler.play_sound") + mocker.patch("basilisk.completion_handler.stop_sound") + handler = CompletionHandler( + on_completion_end=on_completion_end, on_stream_chunk=stream_chunks + ) + + handler.start_completion( + engine=old_engine, + system_message=None, + conversation=empty_conversation, + new_block=message_block, + stream=True, + ) + assert old_stream.read_started.wait(timeout=1) + handler.stop_completion() + assert message_block.response.content == "" + + handler.start_completion( + engine=new_engine, + system_message=None, + conversation=empty_conversation, + new_block=message_block, + stream=True, + ) + assert new_completion_finished.wait(timeout=1) + assert message_block.response.content == "new response." + stream_chunks.assert_called_once_with("new response.") + assert completion_results == [False, True] + assert not handler._stream_buffers + + +def test_error_stream_partial_buffer_does_not_leak_to_restart( + mocker, empty_conversation, message_block +): + """An errored stream discards its partial text before the next request.""" + partial_processed = threading.Event() + release_error = threading.Event() + + def stream_with_error(): + yield "old partial" + partial_processed.set() + assert release_error.wait(timeout=1) + raise RuntimeError("stream failed") + + old_engine = MagicMock() + old_engine.completion.return_value = stream_with_error() + old_engine.completion_response_with_stream.side_effect = lambda response: ( + response + ) + new_engine = MagicMock() + new_engine.completion.return_value = iter(["new response."]) + new_engine.completion_response_with_stream.side_effect = lambda response: ( + response + ) + stream_chunks = MagicMock() + on_error = MagicMock() + new_completion_finished = threading.Event() + + def on_completion_end(success): + if success: + new_completion_finished.set() + + mocker.patch( + "basilisk.completion_handler.wx.CallAfter", + side_effect=lambda callback, *args: callback(*args), + ) + mocker.patch("basilisk.completion_handler.play_sound") + mocker.patch("basilisk.completion_handler.stop_sound") + handler = CompletionHandler( + on_completion_end=on_completion_end, + on_error=on_error, + on_stream_chunk=stream_chunks, + ) + + handler.start_completion( + engine=old_engine, + system_message=None, + conversation=empty_conversation, + new_block=message_block, + stream=True, + ) + old_task = handler.task + assert old_task is not None + assert partial_processed.wait(timeout=1) + release_error.set() + old_task.join(timeout=1) + assert not old_task.is_alive() + on_error.assert_called_once_with("stream failed") + assert message_block.response.content == "" + + handler.start_completion( + engine=new_engine, + system_message=None, + conversation=empty_conversation, + new_block=message_block, + stream=True, + ) + assert new_completion_finished.wait(timeout=1) + assert message_block.response.content == "new response." + stream_chunks.assert_called_once_with("new response.") + assert not handler._stream_buffers + + +def test_stale_stop_callback_does_not_fail_replacement_completion( + mocker, empty_conversation, message_block +): + """A queued stop notification belongs only to the stopped request.""" + queued_callbacks = [] + old_stream = _BlockingStream() + old_engine = MagicMock() + old_engine.completion.return_value = old_stream + old_engine.completion_response_with_stream.return_value = old_stream + old_engine.cancel_completion.side_effect = lambda response: response.close() + second_started = threading.Event() + allow_second_completion = threading.Event() + second_engine = MagicMock() + + def start_second_completion(**kwargs): + second_started.set() + assert allow_second_completion.wait(timeout=1) + return object() + + second_engine.completion.side_effect = start_second_completion + second_engine.completion_response_without_stream.side_effect = ( + lambda **kwargs: kwargs["new_block"] + ) + completion_end = MagicMock() + mocker.patch( + "basilisk.completion_handler.wx.CallAfter", + side_effect=lambda callback, *args: queued_callbacks.append( + (callback, args) + ), + ) + mocker.patch("basilisk.completion_handler.play_sound") + mocker.patch("basilisk.completion_handler.stop_sound") + handler = CompletionHandler(on_completion_end=completion_end) + + handler.start_completion( + engine=old_engine, + system_message=None, + conversation=empty_conversation, + new_block=message_block, + stream=True, + ) + assert old_stream.read_started.wait(timeout=1) + handler.stop_completion() + assert len(queued_callbacks) == 1 + + handler.start_completion( + engine=second_engine, + system_message=None, + conversation=empty_conversation, + new_block=message_block, + ) + assert second_started.wait(timeout=1) + second_task = handler.task + assert second_task is not None + + callback, args = queued_callbacks.pop(0) + callback(*args) + completion_end.assert_not_called() + assert handler.task is second_task + + allow_second_completion.set() + second_task.join(timeout=1) + callback, args = queued_callbacks.pop(0) + callback(*args) + completion_end.assert_called_once_with(True)