diff --git a/docs/chatbook.md b/docs/chatbook.md index ff5a3ebc..c902c4d6 100644 --- a/docs/chatbook.md +++ b/docs/chatbook.md @@ -8,6 +8,8 @@ Generation uses Notebook Intelligence (`POST /notebook-intelligence/chatbook/gen The Chatbook setting **Execution kernel** lists installed Jupyter kernelspecs except `chatbook`. The default is `python3` when that spec exists, otherwise the first Python spec, otherwise the first non-chatbook spec. Change requires restarting open Chatbook notebooks so the child kernel is recreated. +Kernelspecs are listed and started through the Jupyter server's kernelspec manager, so a custom manager works unchanged. With [nb_conda_kernels](https://github.com/anaconda/nb_conda_kernels), for example, the Chatbook kernel is `Chatbook [conda env:base]` (`conda-base-chatbook`) and the backends are its names such as `conda-base-py`. Chatbook is recognized by its kernelspec language, `chatbook`, whatever the kernelspec is named. + Tab completion and contextual help are proxied to that child kernel. Interrupt uses Jupyter's message-mode interrupt so it reaches the child rather than only the Chatbook wrapper. Cell badges show **NL** (natural language) and **Cd** (code). Code cells use the backend language for highlighting. diff --git a/notebook_intelligence/chatbook_kernel/backend.py b/notebook_intelligence/chatbook_kernel/backend.py index 18acd127..e71b87d4 100644 --- a/notebook_intelligence/chatbook_kernel/backend.py +++ b/notebook_intelligence/chatbook_kernel/backend.py @@ -4,6 +4,7 @@ from __future__ import annotations +import json import logging import os import time @@ -13,8 +14,13 @@ log = logging.getLogger(__name__) CHATBOOK_KERNEL_NAME = "chatbook" +CHATBOOK_LANGUAGE = "chatbook" DEFAULT_BACKEND_KERNEL_NAME = "python3" +# The Jupyter server's kernelspec manager, described for the Chatbook kernel +# process (see `describe_kernel_spec_manager`). +KERNEL_SPEC_MANAGER_ENV = "NBI_CHATBOOK_KERNEL_SPEC_MANAGER" + _IOPUB_RELAY_TYPES = { "stream", "display_data", @@ -48,11 +54,23 @@ def _spec_fields(record: Any) -> tuple[str, str]: return language, display_name +def is_chatbook_spec(record: Any) -> bool: + """Whether a kernelspec runs the Chatbook kernel, whatever it is named. + + A custom kernelspec manager can list Chatbook under its own name + (nb_conda_kernels lists it as ``conda-base-chatbook``), so the name alone + does not identify it; the language does. Accepts a ``KernelSpec``, a spec + dict, or a ``get_all_specs()`` record. + """ + language, _ = _spec_fields(record) + return language.lower() == CHATBOOK_LANGUAGE + + def list_backend_kernels(specs: Optional[dict] = None) -> list[dict[str, str]]: """Installed kernelspecs excluding the Chatbook wrapper itself.""" backends: list[dict[str, str]] = [] for name, record in (specs or {}).items(): - if not name or name == CHATBOOK_KERNEL_NAME: + if not name or name == CHATBOOK_KERNEL_NAME or is_chatbook_spec(record): continue language, display_name = _spec_fields(record) backends.append( @@ -66,13 +84,98 @@ def list_backend_kernels(specs: Optional[dict] = None) -> list[dict[str, str]]: return backends +_kernel_spec_manager: Any = None + + +def describe_kernel_spec_manager(manager: Any) -> str: + """The class and JSON-safe config of ``manager``, for the Chatbook kernel. + + The Jupyter server may use a custom kernelspec manager that lists kernels + under its own names (nb_conda_kernels lists ``python3`` as + ``conda-base-py``). The Chatbook kernel has to resolve backend names the + same way, so the server passes this description to it in + ``KERNEL_SPEC_MANAGER_ENV``. Only the config sections of the manager's own + classes are kept. Sets (``allowed_kernelspecs``) travel as lists, which + traitlets turns back into sets; other values that are not JSON are dropped. + """ + cls = type(manager) + config = getattr(manager, "config", None) or {} + sections: dict[str, dict] = {} + for klass in cls.__mro__: + if klass.__name__ not in config: + continue + section = {} + for key, value in dict(config[klass.__name__]).items(): + if isinstance(value, (set, frozenset)): + value = sorted(value, key=str) + try: + json.dumps(value) + except (TypeError, ValueError): + continue + section[key] = value + if section: + sections[klass.__name__] = section + return json.dumps( + {"class": f"{cls.__module__}.{cls.__qualname__}", "config": sections} + ) + + +def set_kernel_spec_manager(manager: Any) -> None: + """Resolve backend kernels with ``manager`` in this process. + + The server extension passes the Jupyter server's own manager, so names + resolved in the server match the ones the kernel picker lists. + """ + global _kernel_spec_manager + _kernel_spec_manager = manager + + +def _kernel_spec_manager_from_env() -> Any: + raw = os.environ.get(KERNEL_SPEC_MANAGER_ENV, "").strip() + if not raw: + return None + try: + from traitlets.config import Config + from traitlets.utils.importstring import import_item + + description = json.loads(raw) + cls = import_item(str(description["class"])) + return cls(config=Config(description.get("config") or {})) + except Exception as exc: + log.warning( + "Could not create the Jupyter server's kernelspec manager (%s); " + "using Jupyter's default", + exc, + ) + return None + + +def kernel_spec_manager() -> Any: + """The kernelspec manager backend kernels are listed and started with. + + In the server, the server's own (see `set_kernel_spec_manager`). In the + Chatbook kernel, one built from the description the server put in the + kernel's environment, so a backend named in Settings → Chatbook resolves + to the same kernelspec here. Otherwise Jupyter's default manager. + """ + global _kernel_spec_manager + if _kernel_spec_manager is None: + manager = _kernel_spec_manager_from_env() + if manager is None: + from jupyter_client.kernelspec import KernelSpecManager + + manager = KernelSpecManager() + _kernel_spec_manager = manager + return _kernel_spec_manager + + def load_kernel_specs() -> dict: try: - from jupyter_client.kernelspec import KernelSpecManager + manager = kernel_spec_manager() except ImportError: return {} try: - return KernelSpecManager().get_all_specs() + return manager.get_all_specs() except Exception as exc: log.warning("Could not list Jupyter kernelspecs: %s", exc) return {} @@ -91,7 +194,7 @@ def resolve_backend_kernel( ) by_name = {item["name"]: item for item in backends} wanted = (preferred or "").strip() - if wanted == CHATBOOK_KERNEL_NAME: + if wanted == CHATBOOK_KERNEL_NAME or is_chatbook_spec((specs or {}).get(wanted)): wanted = "" if wanted in by_name: return by_name[wanted] @@ -163,7 +266,9 @@ def start(self) -> None: if factory is None: from jupyter_client import KernelManager - factory = lambda name: KernelManager(kernel_name=name) + factory = lambda name: KernelManager( + kernel_name=name, kernel_spec_manager=kernel_spec_manager() + ) self._km = factory(self.kernel_name) start_kernel = getattr(self._km, "start_kernel", None) if callable(start_kernel): diff --git a/notebook_intelligence/extension.py b/notebook_intelligence/extension.py index f1241133..332793a9 100644 --- a/notebook_intelligence/extension.py +++ b/notebook_intelligence/extension.py @@ -111,6 +111,12 @@ from notebook_intelligence.rule_injector import has_chatbook_guidelines from notebook_intelligence.chatbook_mentions import list_chatbook_mentions from notebook_intelligence.chatbook_kernel.codegen import ChatbookCodegenError +from notebook_intelligence.chatbook_kernel.backend import ( + KERNEL_SPEC_MANAGER_ENV, + describe_kernel_spec_manager, + is_chatbook_spec, + set_kernel_spec_manager, +) ai_service_manager: AIServiceManager = None log = logging.getLogger(__name__) @@ -203,6 +209,10 @@ def _set_chatbook_kernelspec_execution_cap(kernel_spec_manager, max_mode: str) - The wrapper kernel cannot read the server extension's traitlets directly. Adding the resolved value to its live kernelspec makes env- and traitlet-configured caps equivalent at the actual execution decision. + The same env carries a description of this kernelspec manager, so the + kernel resolves backend kernels under the names this server lists them by. + Chatbook is matched by language, since a custom manager such as + nb_conda_kernels renames its kernelspec. """ if kernel_spec_manager is None: return @@ -211,14 +221,16 @@ def _set_chatbook_kernelspec_execution_cap(kernel_spec_manager, max_mode: str) - if getattr(kernel_spec_manager, "_nbi_chatbook_cap_wrapped", False): return orig_get = kernel_spec_manager.get_kernel_spec + manager_description = describe_kernel_spec_manager(kernel_spec_manager) def get_kernel_spec(kernel_name, *args, **kwargs): spec = orig_get(kernel_name, *args, **kwargs) - if kernel_name == CHATBOOK_KERNEL_NAME: + if kernel_name == CHATBOOK_KERNEL_NAME or is_chatbook_spec(spec): env = dict(getattr(spec, "env", None) or {}) env["NBI_CHATBOOK_MAX_EXECUTION_MODE"] = ( kernel_spec_manager._nbi_chatbook_max_execution_mode ) + env[KERNEL_SPEC_MANAGER_ENV] = manager_description spec.env = env return spec @@ -240,6 +252,8 @@ def _hide_chatbook_kernelspec(kernel_spec_manager) -> None: The kernelspec stays on disk from the package data files; filtering the running manager is what removes the launcher tile and kernel picker. + Chatbook is matched by language as well as name, so a copy a custom + manager renamed (nb_conda_kernels' ``conda-base-chatbook``) goes too. Idempotent so a second call (tests, reload) does not wrap twice. """ if kernel_spec_manager is None: @@ -252,15 +266,25 @@ def _hide_chatbook_kernelspec(kernel_spec_manager) -> None: orig_get = kernel_spec_manager.get_kernel_spec orig_all = getattr(kernel_spec_manager, "get_all_specs", None) + def is_chatbook(kernel_name) -> bool: + if kernel_name == CHATBOOK_KERNEL_NAME: + return True + try: + return is_chatbook_spec(orig_get(kernel_name)) + except Exception: + return False + def find_kernel_specs(*args, **kwargs): specs = orig_find(*args, **kwargs) - specs.pop(CHATBOOK_KERNEL_NAME, None) - return specs + return {name: path for name, path in specs.items() if not is_chatbook(name)} def get_kernel_spec(kernel_name, *args, **kwargs): if kernel_name == CHATBOOK_KERNEL_NAME: raise NoSuchKernel(kernel_name) - return orig_get(kernel_name, *args, **kwargs) + spec = orig_get(kernel_name, *args, **kwargs) + if is_chatbook_spec(spec): + raise NoSuchKernel(kernel_name) + return spec kernel_spec_manager.find_kernel_specs = find_kernel_specs kernel_spec_manager.get_kernel_spec = get_kernel_spec @@ -268,8 +292,11 @@ def get_kernel_spec(kernel_name, *args, **kwargs): def get_all_specs(*args, **kwargs): specs = orig_all(*args, **kwargs) - specs.pop(CHATBOOK_KERNEL_NAME, None) - return specs + return { + name: record + for name, record in specs.items() + if name != CHATBOOK_KERNEL_NAME and not is_chatbook_spec(record) + } kernel_spec_manager.get_all_specs = get_all_specs kernel_spec_manager._nbi_chatbook_hidden = True @@ -4829,6 +4856,11 @@ def _setup_handlers(self, web_app, feature_policies: dict, string_overrides: dic FileUploadHandler.upload_retention_hours = _resolve_positive_int_with_env( "NBI_UPLOAD_RETENTION_HOURS", self.upload_retention_hours ) + # Resolve Chatbook backends in this process the way this server lists + # kernelspecs (a custom manager such as nb_conda_kernels renames them). + kernel_spec_manager = getattr(self.serverapp, "kernel_spec_manager", None) + if kernel_spec_manager is not None: + set_kernel_spec_manager(kernel_spec_manager) self._publish_policies(feature_policies, string_overrides) NotebookIntelligence.handlers = [ (route_pattern_capabilities, GetCapabilitiesHandler), diff --git a/src/chatbook-core.ts b/src/chatbook-core.ts index 0eb7d934..59739696 100644 --- a/src/chatbook-core.ts +++ b/src/chatbook-core.ts @@ -169,8 +169,38 @@ export interface IChatbookExecuteMeta { approvedCode?: string; } +export function isChatbookLanguage( + language: string | undefined | null +): boolean { + return (language ?? '').trim().toLowerCase() === CHATBOOK_LANGUAGE; +} + +/** + * Names of the installed kernelspecs that run the Chatbook kernel. A custom + * kernelspec manager can list Chatbook under its own name (nb_conda_kernels + * lists it as `conda-base-chatbook`), so the name alone does not identify it; + * the kernelspec's language does. + */ +const chatbookKernelNames = new Set([CHATBOOK_KERNEL_NAME]); + +/** Record which installed kernelspecs are Chatbook's, by language. */ +export function setChatbookKernelSpecs( + kernelspecs: + | Record + | null + | undefined +): void { + chatbookKernelNames.clear(); + chatbookKernelNames.add(CHATBOOK_KERNEL_NAME); + for (const [name, spec] of Object.entries(kernelspecs ?? {})) { + if (isChatbookLanguage(spec?.language)) { + chatbookKernelNames.add(name); + } + } +} + export function isChatbookKernelName(name: string | undefined | null): boolean { - return (name ?? '').trim() === CHATBOOK_KERNEL_NAME; + return chatbookKernelNames.has((name ?? '').trim()); } export function isChatbookPromptInlineCompletion( diff --git a/src/chatbook.ts b/src/chatbook.ts index b018de0c..29780373 100644 --- a/src/chatbook.ts +++ b/src/chatbook.ts @@ -57,7 +57,8 @@ export { getChatbookCellMeta, getChatbookCellMode, isChatbookKernelName, - isChatbookPromptInlineCompletion + isChatbookPromptInlineCompletion, + setChatbookKernelSpecs } from './chatbook-core'; let codeCellExecutePatched = false; diff --git a/src/index.ts b/src/index.ts index 8c638844..53dfd06f 100644 --- a/src/index.ts +++ b/src/index.ts @@ -151,6 +151,7 @@ import { attachTerminalDragDrop } from './terminal-drag'; import { DEFAULT_NOTEBOOK_KERNEL, NotebookKernelNotFoundError, + chatbookKernelProfile, findKernelProfile, listKernelProfiles, normalizeNotebookLanguage @@ -158,17 +159,18 @@ import { import { CommandIDs } from './command-ids'; import { - CHATBOOK_KERNEL_NAME, CHATBOOK_LANGUAGE, attachChatbookNotebooks, getChatbookBackendLanguage, getChatbookCellMode, getChatbookCellMeta, + isChatbookKernelName, isChatbookPromptInlineCompletion, isChatbookSession, nextChatbookNotebookMode, patchCodeCellExecute, registerChatbookLanguage, + setChatbookKernelSpecs, summarizeCodeCell, toggleActiveChatbookCellMode, toggleAllChatbookCellModes @@ -1000,6 +1002,13 @@ const plugin: JupyterFrontEndPlugin = { await NBIAPI.initialize(); if (NBIAPI.config.chatbookEnabled) { + // Recognize Chatbook's kernelspec under any name a kernelspec manager + // gives it (nb_conda_kernels lists it as `conda-base-chatbook`). + const kernelspecs = app.serviceManager.kernelspecs; + const recordChatbookKernelSpecs = () => + setChatbookKernelSpecs(kernelspecs.specs?.kernelspecs); + void kernelspecs.ready.then(recordChatbookKernelSpecs); + kernelspecs.specsChanged.connect(recordChatbookKernelSpecs); registerChatbookLanguage(languageRegistry); patchCodeCellExecute(); attachChatbookNotebooks(notebookTracker, { @@ -1477,8 +1486,11 @@ const plugin: JupyterFrontEndPlugin = { icon: () => sidebarIcon, isVisible: () => NBIAPI.config.chatbookEnabled, execute: async () => { + const kernelspecs = app.serviceManager.kernelspecs; + await kernelspecs.ready; return app.commands.execute(CommandIDs.createNewNotebook, { - kernelName: CHATBOOK_KERNEL_NAME + kernelName: chatbookKernelProfile(kernelspecs.specs?.kernelspecs) + .kernelName }); } }); @@ -1492,7 +1504,7 @@ const plugin: JupyterFrontEndPlugin = { return false; } if ( - current.sessionContext.session?.kernel?.name !== CHATBOOK_KERNEL_NAME + !isChatbookKernelName(current.sessionContext.session?.kernel?.name) ) { return false; } diff --git a/src/notebook-kernels.ts b/src/notebook-kernels.ts index 0a84efcb..9c1b3e95 100644 --- a/src/notebook-kernels.ts +++ b/src/notebook-kernels.ts @@ -2,7 +2,7 @@ import { KernelSpec, KernelSpecManager } from '@jupyterlab/services'; -import { chatbookLanguageId } from './chatbook-core'; +import { chatbookLanguageId, isChatbookLanguage } from './chatbook-core'; export interface INotebookKernelProfile { language: string; @@ -113,11 +113,33 @@ export function listKernelProfiles( }); } +/** + * The kernelspec a new Chatbook opens with: `chatbook` when it is installed + * under that name, otherwise the first kernelspec in Chatbook's language (a + * custom kernelspec manager such as nb_conda_kernels renames it). + */ +export function chatbookKernelProfile( + specs: Record | undefined +): INotebookKernelProfile { + const profiles = listKernelProfiles(specs).filter(profile => + isChatbookLanguage(profile.language) + ); + return ( + profiles.find( + profile => profile.kernelName === CHATBOOK_KERNEL.kernelName + ) ?? + profiles[0] ?? + CHATBOOK_KERNEL + ); +} + export function listChatbookBackendProfiles( specs: Record | undefined ): INotebookKernelProfile[] { return listKernelProfiles(specs).filter( - profile => profile.kernelName !== CHATBOOK_KERNEL.kernelName + profile => + profile.kernelName !== CHATBOOK_KERNEL.kernelName && + !isChatbookLanguage(profile.language) ); } diff --git a/tests/conftest.py b/tests/conftest.py index 1cbb142a..87b8f625 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -6,6 +6,9 @@ from pathlib import Path from unittest.mock import MagicMock, patch +from jupyter_client.kernelspec import KernelSpec, KernelSpecManager, NoSuchKernel +from traitlets import Unicode + from notebook_intelligence.ruleset import RuleContext from notebook_intelligence.config import NBIConfig @@ -232,4 +235,31 @@ def invalid_rules_directory(temp_rules_directory): with open(rules_path / "no-frontmatter.md", 'w') as f: f.write(no_frontmatter) - return temp_rules_directory \ No newline at end of file + return temp_rules_directory + +class RenamingKernelSpecManager(KernelSpecManager): + """Lists kernelspecs under prefixed names, the way nb_conda_kernels does. + + nb_conda_kernels lists ``chatbook`` as ``conda-base-chatbook`` and + ``python3`` as ``conda-base-py``. ``prefix`` is configurable so a test can + tell a manager built from the server's config from a default one. + """ + + prefix = Unicode("conda-base-").tag(config=True) + + _LANGUAGES = {"chatbook": "chatbook", "py": "python"} + + def find_kernel_specs(self): + return {self.prefix + name: f"/kernels/{name}" for name in self._LANGUAGES} + + def get_kernel_spec(self, kernel_name, *args, **kwargs): + for name, language in self._LANGUAGES.items(): + if kernel_name == self.prefix + name: + return KernelSpec( + name=kernel_name, + display_name=f"{name} [conda env:base]", + language=language, + argv=["python", "-m", "ipykernel_launcher", "-f", "{connection_file}"], + resource_dir=f"/kernels/{name}", + ) + raise NoSuchKernel(kernel_name) diff --git a/tests/test_chatbook.py b/tests/test_chatbook.py index 9a09ee5f..9ab8d4c8 100644 --- a/tests/test_chatbook.py +++ b/tests/test_chatbook.py @@ -1491,6 +1491,22 @@ class Host: backend.load_kernel_specs = original +def test_backend_kernel_name_uses_the_servers_kernel_spec_manager(monkeypatch): + # Under nb_conda_kernels the server lists `python3` as `conda-base-py`, and + # rules have to be matched against the name the user sees. + from notebook_intelligence import chatbook_generate + from notebook_intelligence.chatbook_kernel import backend + from tests.conftest import RenamingKernelSpecManager + + monkeypatch.setattr(backend, '_kernel_spec_manager', None) + backend.set_kernel_spec_manager(RenamingKernelSpecManager()) + + class Host: + nbi_config = type('Cfg', (), {'chatbook_backend_kernel': ''})() + + assert chatbook_generate.chatbook_backend_kernel_name(Host()) == 'conda-base-py' + + def test_backend_kernel_name_ignores_a_config_naming_the_wrapper(): # The wrapper cannot be its own backend. A hand-edited config can still # name it, and matching rules against it would reinstate the original bug. diff --git a/tests/test_chatbook_backend.py b/tests/test_chatbook_backend.py index ca4a8828..0c61dfb7 100644 --- a/tests/test_chatbook_backend.py +++ b/tests/test_chatbook_backend.py @@ -6,10 +6,20 @@ import pytest from jupyter_client.session import Session +from jupyter_client.kernelspec import KernelSpecManager +from traitlets.config import Config + +from notebook_intelligence.chatbook_kernel import backend as backend_module from notebook_intelligence.chatbook_kernel.backend import ( + KERNEL_SPEC_MANAGER_ENV, ChatbookBackend, + describe_kernel_spec_manager, + is_chatbook_spec, + kernel_spec_manager, list_backend_kernels, + load_kernel_specs, resolve_backend_kernel, + set_kernel_spec_manager, ) from notebook_intelligence.chatbook_kernel.codegen import ( cell_codegen_instructions, @@ -18,6 +28,7 @@ from notebook_intelligence.chatbook_kernel.danger import scan_generated_code from notebook_intelligence.chatbook_kernel.nbi_client import NBIClientError from notebook_intelligence.chatbook_kernel.kernel import ChatbookKernel, is_code_execute +from tests.conftest import RenamingKernelSpecManager SPECS = { @@ -52,6 +63,135 @@ def test_resolve_backend_kernel_missing_name_raises(): resolve_backend_kernel('does-not-exist', SPECS) +# How nb_conda_kernels lists the same kernels: Chatbook keeps its language +# under a new name, and there is no `python3`. +CONDA_SPECS = { + 'conda-base-chatbook': { + 'spec': {'language': 'chatbook', 'display_name': 'Chatbook [conda env:base] *'} + }, + 'conda-base-py': { + 'spec': {'language': 'python', 'display_name': 'Python [conda env:base] *'} + }, + 'conda-base-ir': {'spec': {'language': 'R', 'display_name': 'R [conda env:base] *'}}, +} + + +@pytest.fixture +def fresh_kernel_spec_manager(monkeypatch): + """Forget the process-wide kernelspec manager and any server description.""" + monkeypatch.setattr(backend_module, '_kernel_spec_manager', None) + monkeypatch.delenv(KERNEL_SPEC_MANAGER_ENV, raising=False) + + +def test_is_chatbook_spec_goes_by_language_not_name(): + assert is_chatbook_spec(CONDA_SPECS['conda-base-chatbook']) + assert is_chatbook_spec(SimpleNamespace(language='Chatbook')) + assert not is_chatbook_spec(CONDA_SPECS['conda-base-py']) + assert not is_chatbook_spec(None) + + +def test_list_backend_kernels_excludes_a_renamed_chatbook(): + names = [item['name'] for item in list_backend_kernels(CONDA_SPECS)] + assert names == ['conda-base-ir', 'conda-base-py'] + + +def test_resolve_backend_kernel_accepts_renamed_kernels(): + assert resolve_backend_kernel('conda-base-py', CONDA_SPECS)['name'] == 'conda-base-py' + # With no `python3`, the default is the first Python kernel. + assert resolve_backend_kernel('', CONDA_SPECS)['name'] == 'conda-base-py' + # A renamed Chatbook cannot be its own backend any more than `chatbook` can. + assert ( + resolve_backend_kernel('conda-base-chatbook', CONDA_SPECS)['name'] + == 'conda-base-py' + ) + + +def test_kernel_lists_backends_with_the_servers_kernel_spec_manager( + monkeypatch, fresh_kernel_spec_manager +): + allowed = {'conda-env-tools-chatbook', 'conda-env-tools-py'} + # The server describes its manager, config included, in the kernel's env ... + server_manager = RenamingKernelSpecManager( + config=Config( + { + 'RenamingKernelSpecManager': {'prefix': 'conda-env-tools-'}, + 'KernelSpecManager': {'allowed_kernelspecs': allowed}, + } + ) + ) + monkeypatch.setenv( + KERNEL_SPEC_MANAGER_ENV, describe_kernel_spec_manager(server_manager) + ) + + # ... so the kernel resolves the names Settings and the kernel picker show. + manager = kernel_spec_manager() + assert type(manager) is RenamingKernelSpecManager + assert manager.prefix == 'conda-env-tools-' + assert manager.allowed_kernelspecs == allowed + specs = load_kernel_specs() + assert sorted(specs) == ['conda-env-tools-chatbook', 'conda-env-tools-py'] + assert resolve_backend_kernel('conda-env-tools-py', specs)['language'] == 'python' + + +def test_describe_kernel_spec_manager_keeps_only_its_own_json_config(): + manager = RenamingKernelSpecManager( + config=Config( + { + 'RenamingKernelSpecManager': {'prefix': 'p-', 'not_json': object()}, + 'ServerApp': {'port': 8888}, + } + ) + ) + description = json.loads(describe_kernel_spec_manager(manager)) + assert description['class'] == ( + f'{RenamingKernelSpecManager.__module__}.RenamingKernelSpecManager' + ) + assert description['config'] == {'RenamingKernelSpecManager': {'prefix': 'p-'}} + + +def test_kernel_spec_manager_defaults_without_a_server_description( + fresh_kernel_spec_manager, +): + assert type(kernel_spec_manager()) is KernelSpecManager + + +def test_kernel_spec_manager_falls_back_on_a_bad_description( + monkeypatch, fresh_kernel_spec_manager +): + monkeypatch.setenv(KERNEL_SPEC_MANAGER_ENV, '{"class": "no_such_module.Manager"}') + assert type(kernel_spec_manager()) is KernelSpecManager + + +def test_server_resolves_backends_with_its_own_kernel_spec_manager( + fresh_kernel_spec_manager, +): + server_manager = RenamingKernelSpecManager() + set_kernel_spec_manager(server_manager) + assert kernel_spec_manager() is server_manager + assert resolve_backend_kernel('', load_kernel_specs())['name'] == 'conda-base-py' + + +def test_backend_starts_through_the_resolved_kernel_spec_manager( + monkeypatch, fresh_kernel_spec_manager +): + import jupyter_client + + server_manager = RenamingKernelSpecManager() + set_kernel_spec_manager(server_manager) + created = {} + + def kernel_manager(**kwargs): + created.update(kwargs) + return SimpleNamespace() + + monkeypatch.setattr(jupyter_client, 'KernelManager', kernel_manager) + ChatbookBackend('conda-base-py').start() + assert created == { + 'kernel_name': 'conda-base-py', + 'kernel_spec_manager': server_manager, + } + + def test_is_code_execute_only_accepts_code_mode(): assert is_code_execute({'executeMode': 'code'}) assert not is_code_execute({'executeMode': 'prompt'}) diff --git a/tests/test_chatbook_enable.py b/tests/test_chatbook_enable.py index ae72bcd0..1d797405 100644 --- a/tests/test_chatbook_enable.py +++ b/tests/test_chatbook_enable.py @@ -6,6 +6,7 @@ import pytest from jupyter_client.kernelspec import NoSuchKernel +from notebook_intelligence.chatbook_kernel.backend import KERNEL_SPEC_MANAGER_ENV from notebook_intelligence.extension import ( CHATBOOK_DISABLED_MESSAGE, FEATURE_POLICY_DEFAULTS, @@ -25,6 +26,7 @@ POLICY_USER_CHOICE, is_force_off, ) +from tests.conftest import RenamingKernelSpecManager def test_chatbook_enabled_defaults_on(): @@ -105,6 +107,29 @@ def test_chatbook_kernelspec_inherits_resolved_traitlet_cap(): assert spec.env["NBI_CHATBOOK_MAX_EXECUTION_MODE"] == "confirm-if-risky" +def test_hide_chatbook_kernelspec_drops_a_renamed_chatbook(): + # nb_conda_kernels lists Chatbook as conda-base-chatbook; force-off has to + # hide that one too. + manager = RenamingKernelSpecManager() + _hide_chatbook_kernelspec(manager) + assert list(manager.find_kernel_specs()) == ["conda-base-py"] + assert list(manager.get_all_specs()) == ["conda-base-py"] + with pytest.raises(NoSuchKernel): + manager.get_kernel_spec("conda-base-chatbook") + assert manager.get_kernel_spec("conda-base-py").name == "conda-base-py" + + +def test_renamed_chatbook_kernelspec_gets_the_cap_and_the_server_manager(): + manager = RenamingKernelSpecManager() + _set_chatbook_kernelspec_execution_cap(manager, "always-confirm") + + env = manager.get_kernel_spec("conda-base-chatbook").env + assert env["NBI_CHATBOOK_MAX_EXECUTION_MODE"] == "always-confirm" + description = json.loads(env[KERNEL_SPEC_MANAGER_ENV]) + assert description["class"].endswith(".RenamingKernelSpecManager") + assert manager.get_kernel_spec("conda-base-py").env == {} + + class _DummyHandler: def __init__(self, chatbook_enabled: bool): self.chatbook_enabled = chatbook_enabled diff --git a/tests/ts/chatbook-kernel-names.test.ts b/tests/ts/chatbook-kernel-names.test.ts new file mode 100644 index 00000000..8fb1a4f2 --- /dev/null +++ b/tests/ts/chatbook-kernel-names.test.ts @@ -0,0 +1,89 @@ +// Copyright (c) Mehmet Bektas + +import { + isChatbookKernelName, + isChatbookPromptInlineCompletion, + setChatbookKernelSpecs +} from '../../src/chatbook-core'; +import { + chatbookKernelProfile, + listChatbookBackendProfiles, + resolveChatbookBackendProfile +} from '../../src/notebook-kernels'; + +// How nb_conda_kernels lists the kernels of a conda environment: Chatbook +// keeps its language under a new name, and there is no `python3`. +const CONDA_SPECS = { + 'conda-base-chatbook': { + name: 'conda-base-chatbook', + language: 'chatbook', + display_name: 'Chatbook [conda env:base] *' + }, + 'conda-base-py': { + name: 'conda-base-py', + language: 'python', + display_name: 'Python [conda env:base] *' + }, + 'conda-base-ir': { + name: 'conda-base-ir', + language: 'R', + display_name: 'R [conda env:base] *' + } +} as any; + +const PLAIN_SPECS = { + chatbook: { + name: 'chatbook', + language: 'chatbook', + display_name: 'Chatbook' + }, + python3: { name: 'python3', language: 'python', display_name: 'Python 3' } +} as any; + +describe('Chatbook kernelspecs under other names', () => { + afterEach(() => { + setChatbookKernelSpecs(undefined); + }); + + it('recognizes a kernelspec in Chatbook language whatever it is named', () => { + expect(isChatbookKernelName('conda-base-chatbook')).toBe(false); + + setChatbookKernelSpecs(CONDA_SPECS); + + expect(isChatbookKernelName('conda-base-chatbook')).toBe(true); + expect(isChatbookKernelName('chatbook')).toBe(true); + expect(isChatbookKernelName('conda-base-py')).toBe(false); + expect(isChatbookPromptInlineCompletion('conda-base-chatbook')).toBe(true); + }); + + it('forgets a renamed Chatbook once it is no longer installed', () => { + setChatbookKernelSpecs(CONDA_SPECS); + setChatbookKernelSpecs(PLAIN_SPECS); + + expect(isChatbookKernelName('conda-base-chatbook')).toBe(false); + expect(isChatbookKernelName('chatbook')).toBe(true); + }); + + it('never offers a renamed Chatbook as its own backend', () => { + expect( + listChatbookBackendProfiles(CONDA_SPECS).map(p => p.kernelName) + ).toEqual(['conda-base-ir', 'conda-base-py']); + expect(resolveChatbookBackendProfile(CONDA_SPECS).kernelName).toBe( + 'conda-base-py' + ); + expect( + resolveChatbookBackendProfile(CONDA_SPECS, 'conda-base-chatbook') + .kernelName + ).toBe('conda-base-py'); + }); + + it('opens a new Chatbook with the installed Chatbook kernelspec', () => { + expect(chatbookKernelProfile(CONDA_SPECS).kernelName).toBe( + 'conda-base-chatbook' + ); + expect( + chatbookKernelProfile({ ...CONDA_SPECS, ...PLAIN_SPECS }).kernelName + ).toBe('chatbook'); + expect(chatbookKernelProfile(undefined).kernelName).toBe('chatbook'); + }); +});