Implement compatibility with nb-conda-kernels

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.

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):

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,24 +266,37 @@ 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
     if orig_all is not None:
 
         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),

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<string>([CHATBOOK_KERNEL_NAME]);
+
+/** Record which installed kernelspecs are Chatbook's, by language. */
+export function setChatbookKernelSpecs(
+  kernelspecs:
+    | Record<string, { language?: string } | undefined>
+    | 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(

src/chatbook.ts

@@ -57,7 +57,8 @@ export {
   getChatbookCellMeta,
   getChatbookCellMode,
   isChatbookKernelName,
-  isChatbookPromptInlineCompletion
+  isChatbookPromptInlineCompletion,
+  setChatbookKernelSpecs
 } from './chatbook-core';
 
 let codeCellExecutePatched = false;

src/index.ts

@@ -151,24 +151,26 @@ import { attachTerminalDragDrop } from './terminal-drag';
 import {
   DEFAULT_NOTEBOOK_KERNEL,
   NotebookKernelNotFoundError,
+  chatbookKernelProfile,
   findKernelProfile,
   listKernelProfiles,
   normalizeNotebookLanguage
 } from './notebook-kernels';
 
 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<INotebookIntelligence> = {
     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<INotebookIntelligence> = {
       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<INotebookIntelligence> = {
           return false;
         }
         if (
-          current.sessionContext.session?.kernel?.name !== CHATBOOK_KERNEL_NAME
+          !isChatbookKernelName(current.sessionContext.session?.kernel?.name)
         ) {
           return false;
         }