fmt, fix

- st_demo: persist failure in call history, clear active_handle,
  and disable Run while call active to avoid overwriting handle
- JupyterApp.shutdown: gather all shutdowns with return_exceptions=True
  and re-raise as BaseExceptionGroup so one failure doesn't leak kernels
- messages: fallback known-but-malformed messages to UnknownJupyterMessage
  via ValidationError, make Unknown channel permissive (str) for nuance
- transport: record reader failures as terminal state (_reader_failure)
  and fail fast from execute/messages_for instead of hanging
- shell: don't reset execution_count to 0 when IPython returns None
- README: use ipython_shell imports

Skipped non-urgent/deferred items: line(), get_shell_2, test renames,
import laziness, call-ID bounding, settle window config (worth doing later).

README.md

@@ -3,8 +3,7 @@
 The application-facing API is:
 
 ```python
-from ipython_demo import App
-from ipython_demo import CallOptions
+from ipython_shell import App, CallOptions
 
 app = App()
 call = app.run_code(

src/ipython_shell/__init__.py

@@ -7,8 +7,8 @@ from .events import (
     event_from_history_output,
 )
 from .jupyter.app import JupyterApp
-from .jupyter.shell import JupyterShell
 from .jupyter.messages import JupyterMessage, KernelMessage, ParsedJupyterMessage
+from .jupyter.shell import JupyterShell
 from .jupyter.transport import InputRequest, JupyterTransport
 from .models import (
     AppInfo,
@@ -46,9 +46,9 @@ __all__ = [
     "InputRequest",
     "JSONValue",
     "JupyterApp",
+    "JupyterMessage",
     "JupyterShell",
     "JupyterTransport",
-    "JupyterMessage",
     "KernelMessage",
     "ParsedJupyterMessage",
     "Shell",

src/ipython_shell/jupyter/__init__.py

@@ -18,17 +18,17 @@ from .shell import JupyterShell
 from .transport import InputRequest, JupyterTransport
 
 __all__ = [
-    "InputRequest",
     "DisplayDataMessage",
     "ErrorMessage",
     "ExecuteReplyMessage",
     "ExecuteResultMessage",
+    "InputRequest",
     "InputRequestMessage",
     "JupyterApp",
     "JupyterMessage",
-    "KernelMessage",
     "JupyterShell",
     "JupyterTransport",
+    "KernelMessage",
     "KnownJupyterMessage",
     "ParsedJupyterMessage",
     "StatusMessage",

src/ipython_shell/jupyter/app.py

@@ -1,3 +1,4 @@
+import asyncio
 from collections.abc import AsyncIterator
 from datetime import UTC, datetime
 from time import monotonic
@@ -97,6 +98,10 @@ class JupyterApp:
 
     async def shutdown(self) -> None:
         """Stop all kernel processes owned by this app."""
-        awaitables = [shell.shutdown() for shell in self.shells.values()]
-        for awaitable in awaitables:
-            await awaitable
+        results = await asyncio.gather(
+            *(shell.shutdown() for shell in self.shells.values()),
+            return_exceptions=True,
+        )
+        errors = [result for result in results if isinstance(result, BaseException)]
+        if errors:
+            raise BaseExceptionGroup("Kernel shutdown failed", errors)

src/ipython_shell/jupyter/messages.py

@@ -1,21 +1,20 @@
 from typing import Annotated, ClassVar, Final, Literal
 
-from pydantic import BaseModel, ConfigDict, Field, TypeAdapter
-
+from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError
 
 type JupyterChannel = Literal["iopub", "shell", "stdin"]
 type MessageStatus = Literal["ok", "error", "abort"]
 
-STREAM: Final[Literal["stream"]] = "stream"
-DISPLAY_DATA: Final[Literal["display_data"]] = "display_data"
-UPDATE_DISPLAY_DATA: Final[Literal["update_display_data"]] = "update_display_data"
-EXECUTE_RESULT: Final[Literal["execute_result"]] = "execute_result"
-ERROR: Final[Literal["error"]] = "error"
-STATUS: Final[Literal["status"]] = "status"
-EXECUTE_INPUT: Final[Literal["execute_input"]] = "execute_input"
-CLEAR_OUTPUT: Final[Literal["clear_output"]] = "clear_output"
-EXECUTE_REPLY: Final[Literal["execute_reply"]] = "execute_reply"
-INPUT_REQUEST: Final[Literal["input_request"]] = "input_request"
+STREAM: Final = "stream"
+DISPLAY_DATA: Final = "display_data"
+UPDATE_DISPLAY_DATA: Final = "update_display_data"
+EXECUTE_RESULT: Final = "execute_result"
+ERROR: Final = "error"
+STATUS: Final = "status"
+EXECUTE_INPUT: Final = "execute_input"
+CLEAR_OUTPUT: Final = "clear_output"
+EXECUTE_REPLY: Final = "execute_reply"
+INPUT_REQUEST: Final = "input_request"
 
 
 class MimeContent(BaseModel):
@@ -176,6 +175,9 @@ class InputRequestMessage(_TypedMessage):
 class UnknownJupyterMessage(JupyterMessage):
     """An extension or future message preserved without lossy parsing."""
 
+    # Allow any channel value so known-but-malformed fallbacks never fail
+    # on a strict Literal. The parser preserves the transport channel verbatim.
+    channel: str = "iopub"  # type: ignore[assignment]
     raw: dict[str, object] = Field(default_factory=dict)
 
 
@@ -252,7 +254,14 @@ def parse_jupyter_message(
 
     if values["msg_type"] not in _KNOWN_MESSAGE_TYPES:
         return UnknownJupyterMessage.model_validate({**values, "raw": raw})
-    return _KNOWN_MESSAGE_ADAPTER.validate_python(values)
+    try:
+        return _KNOWN_MESSAGE_ADAPTER.validate_python(values)
+    except ValidationError:
+        # Known msg_type but content or channel doesn't match the strict
+        # model (e.g., extra execution_state, wrong stdout/stderr name,
+        # or msg delivered on an unexpected channel). Degrade losslessly
+        # instead of aborting the whole call.
+        return UnknownJupyterMessage.model_validate({**values, "raw": raw})
 
 
 # Existing callers use this name for the transport-level message. Keep it as

src/ipython_shell/jupyter/transport.py

@@ -11,8 +11,6 @@ from .messages import (
     ExecuteReplyMessage,
     InputRequestMessage,
     JupyterChannel,
-    JupyterMessage,
-    KernelMessage,
     ParsedJupyterMessage,
     StatusMessage,
     parse_jupyter_message,
@@ -48,6 +46,7 @@ class JupyterTransport:
             asyncio.Queue[tuple[JupyterChannel, dict[str, Any] | BaseException]] | None
         ) = None
         self._reader_tasks: set[asyncio.Task[None]] = set()
+        self._reader_failure: BaseException | None = None
 
     async def start(self) -> None:
         """Start the kernel and its persistent channel readers."""
@@ -100,6 +99,7 @@ class JupyterTransport:
             raise
         except BaseException as error:
             # Surface reader failures instead of leaving messages_for() stuck.
+            self._reader_failure = error
             await self._message_queue.put((channel, error))
 
     async def _stop_readers(self) -> None:
@@ -113,6 +113,10 @@ class JupyterTransport:
 
     async def execute(self, code: str) -> str:
         """Submit one cell and return its Jupyter message ID."""
+        if self._reader_failure is not None:
+            raise RuntimeError(
+                "JupyterTransport channel reader failed"
+            ) from self._reader_failure
         if self.client is None:
             raise RuntimeError("JupyterTransport has not been started")
 
@@ -126,6 +130,10 @@ class JupyterTransport:
         self, call_id: str
     ) -> AsyncIterator[ParsedJupyterMessage | InputRequest]:
         """Yield decoded output, input, and completion messages for one call."""
+        if self._reader_failure is not None:
+            raise RuntimeError(
+                "JupyterTransport channel reader failed"
+            ) from self._reader_failure
         if self.client is None:
             raise RuntimeError("JupyterTransport has not been started")
         if self._active_call != call_id:

src/ipython_shell/shell.py

@@ -87,7 +87,8 @@ class Shell:
         else:
             self.status = "ready"
             self.last_used_at = datetime.now(UTC).isoformat()
-            self._last_execution_count = call_result.execution_count or 0
+            if call_result.execution_count is not None:
+                self._last_execution_count = call_result.execution_count
             return call_result, events
 
     def describe(self) -> ShellInfo:

src/st_demo/app.py

@@ -10,8 +10,7 @@ from pathlib import Path
 from uuid import uuid4
 
 from ipython_shell.jupyter import InputRequest, JupyterApp, ParsedJupyterMessage
-
-from st_demo.ui import event_to_record, preferred_mime
+from st_demo.ui import event_to_record, latex_body, preferred_mime
 
 
 @dataclass
@@ -28,7 +27,9 @@ class _CallFailed:
 
 @dataclass
 class _BackgroundCall:
-    events: queue.Queue[ParsedJupyterMessage | InputRequest | _CallFinished | _CallFailed]
+    events: queue.Queue[
+        ParsedJupyterMessage | InputRequest | _CallFinished | _CallFailed
+    ]
     task: Future[None] | None = None
 
 
@@ -174,7 +175,7 @@ def _render_mime(data: dict[str, object]) -> None:
     elif mime == "application/javascript":
         st.code(str(value), language="javascript")
     elif mime == "text/latex":
-        st.latex(str(value))
+        st.latex(latex_body(value))
     else:
         st.code(str(value), language="text")
 
@@ -227,6 +228,17 @@ def _render_active_call(controller: ShellController) -> None:
         events, finished, failed = controller.poll(handle)
         records = call_records.setdefault(handle, [])
         records.extend(event_to_record(event) for event in events)
+        # Restore persisted failure from a previous poll where the sentinel
+        # was already drained. Without this, the next poll returns
+        # failed=None, finished=False and shows "Running…" forever.
+        if failed is None:
+            failures = st.session_state.get("call_failures")
+            if isinstance(failures, dict):
+                cached = failures.get(handle)
+                if isinstance(cached, BaseException):
+                    failed = cached  # type: ignore[assignment]
+                elif isinstance(cached, Exception):
+                    failed = cached
 
     for record in visible_records(call_records):
         _render_record(record)
@@ -249,7 +261,17 @@ def _render_active_call(controller: ShellController) -> None:
             controller.reply_to_input(handle, call_id, value)
             st.rerun(scope="fragment")
     elif failed is not None:
+        failures = st.session_state.get("call_failures")
+        if not isinstance(failures, dict):
+            failures = {}
+            st.session_state.call_failures = failures
+        failures[handle] = failed
+        # Persist in call history so the error remains visible after clearing.
+        error_text = f"{type(failed).__name__}: {failed}"
+        if not records or records[-1].get("text") != error_text:
+            records.append({"kind": "error", "text": error_text})
         st.exception(failed)
+        st.session_state.active_handle = None
     elif not finished:
         st.caption("Running…")
     else:
@@ -276,7 +298,9 @@ def render() -> None:
         value="import matplotlib.pyplot as plt\nplt.plot([1, 2, 3], [4, 5, 6])",
         height=180,
     )
-    if st.button("Run", type="primary"):
+    active_handle = st.session_state.get("active_handle")
+    is_running = isinstance(active_handle, str)
+    if st.button("Run", type="primary", disabled=is_running):
         handle = controller.start_call(code, shell_name)
         st.session_state.active_handle = handle
         st.session_state.call_records = getattr(

src/st_demo/ui.py

@@ -13,7 +13,6 @@ from ipython_shell.jupyter.messages import (
 )
 from ipython_shell.jupyter.transport import InputRequest
 
-
 MIME_PRIORITY: Final[tuple[str, ...]] = (
     "text/html",
     "image/svg+xml",
@@ -34,6 +33,21 @@ def preferred_mime(data: Mapping[str, object]) -> tuple[str, object] | None:
     return next(iter(data.items()), None)
 
 
+def latex_body(value: object) -> str:
+    """Adapt common Jupyter LaTeX payloads to Streamlit's ``st.latex`` API.
+
+    IPython's LaTeX formatter commonly returns ``$...$`` while Streamlit
+    supplies its own delimiters around the value. Remove only the outer
+    delimiters so they are not rendered as literal dollar signs.
+    """
+    body = str(value).strip()
+    delimiters = (("$$", "$$"), ("$", "$"), (r"\[", r"\]"), (r"\(", r"\)"))
+    for opening, closing in delimiters:
+        if body.startswith(opening) and body.endswith(closing):
+            return body[len(opening) : -len(closing)].strip()
+    return body
+
+
 def _content_record(content: object) -> dict[str, object]:
     if isinstance(content, BaseModel):
         return content.model_dump(mode="json")

tests/test_info.py

@@ -2,10 +2,10 @@ import unittest
 
 from fastapi.testclient import TestClient
 
+from ipython_agent.ipython_wrapper import get_app_info
 from ipython_mcp.app import app as mcp_app
 from ipython_shell import App, AppInfo
 from ipython_webapp.app import app as web_app
-from langchain_demo.ipython_wrapper import get_app_info
 
 
 class AppInfoTests(unittest.IsolatedAsyncioTestCase):

tests/test_jupyter_app.py

@@ -3,8 +3,8 @@ import re
 import unittest
 
 from ipython_shell.jupyter.app import JupyterApp
-from ipython_shell.jupyter.shell import JupyterShell
 from ipython_shell.jupyter.messages import ExecuteResultMessage
+from ipython_shell.jupyter.shell import JupyterShell
 from ipython_shell.jupyter.transport import InputRequest
 
 

tests/test_st_demo.py

@@ -8,7 +8,7 @@ from ipython_shell.jupyter.messages import (
     UnknownJupyterMessage,
 )
 from st_demo.app import visible_records
-from st_demo.ui import event_to_record, preferred_mime
+from st_demo.ui import event_to_record, latex_body, preferred_mime
 
 
 class UIDisplayTests(unittest.TestCase):
@@ -44,6 +44,24 @@ class UIDisplayTests(unittest.TestCase):
 
         self.assertEqual((mime, value), ("text/html", "<b>rich</b>"))
 
+    def test_preferred_mime_chooses_latex_before_plain_text(self):
+        latex = r"\\int_0^1 x^2\\,dx = \\frac{1}{3}"
+
+        mime, value = preferred_mime(
+            {
+                "text/plain": "integral",
+                "text/latex": latex,
+            }
+        )
+
+        self.assertEqual((mime, value), ("text/latex", latex))
+
+    def test_latex_body_unwraps_ipython_math_delimiters(self):
+        self.assertEqual(
+            latex_body(r"$\displaystyle x^{2} - x - 6$"),
+            r"\displaystyle x^{2} - x - 6",
+        )
+
     def test_execute_result_record_keeps_mime_and_repr(self):
         event = ExecuteResultMessage(
             message_id="message-1",

tests/test_transport.py

@@ -211,12 +211,10 @@ class AsyncTransportTests(unittest.IsolatedAsyncioTestCase):
             await transport.start()
             try:
                 call_id = await transport.execute(
-                    "import matplotlib.pyplot as plt\n"
-                    "plt.plot([1, 2, 3], [4, 5, 6])"
+                    "import matplotlib.pyplot as plt\nplt.plot([1, 2, 3], [4, 5, 6])"
                 )
                 messages = [
-                    message
-                    async for message in transport.messages_for(call_id)
+                    message async for message in transport.messages_for(call_id)
                 ]
             finally:
                 await transport.shutdown()