From 700e20e52a914119f8c4628f6365abc01636032f Mon Sep 17 00:00:00 2001
From: Mohamed Elkholy <moelkholy1995@gmail.com>
Date: Sat, 20 Jun 2026 12:49:31 -0400
Subject: [PATCH 1/2] fix(update): stage native updates on exit + move notice
 below input
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Two separate update-flow bugs reported together.

1. The "Updated → vX. Restart to apply." notice rendered inside the prompt
   input box (between the input line and the footer separator). It now renders
   below the separator rule, underneath the box. Helper renamed
   `_prepend_update_notice` → `_append_update_notice` and both toolbar call
   sites place it after the separator.

2. Logging into ChatGPT after a silent auto-update crashed with
   `zlib.error: incorrect header check`. Root cause: the silent updater's
   `_install_native_archive` did `os.replace` over `sys.executable`, overwriting
   the running PyInstaller onefile bundle in place. This build reads its Python
   archive lazily from the exe path, so the first not-yet-loaded import after the
   swap (llm.py `openai_codex` branch) read a stale archive and died.

   The native update now stages the new binary beside the running exe
   (`.{exe}.staged`) and promotes it via `os.replace` at process exit (atexit),
   so it goes live on the next launch — matching the restart notice. The smoke
   check validates the staged binary; a binary that fails smoke is discarded,
   never promoted. A boundary guard converts any residual post-update archive
   corruption into a clean "restart to apply" message instead of a fatal
   traceback.

Verified: make check-pythinker-code (ruff + pyright + ty) green; full test
suite 6388 passed, 7 skipped, 1 xfailed.
---
 CHANGELOG.md                                  | 10 +++
 src/pythinker_code/cli/__init__.py            | 41 +++++++++++
 src/pythinker_code/ui/shell/__init__.py       |  2 +-
 src/pythinker_code/ui/shell/prompt.py         | 15 ++--
 src/pythinker_code/ui/shell/update.py         | 73 +++++++++++++++++--
 .../ui/shell/update_orchestrator.py           | 32 +++++++-
 .../cli/test_post_update_corruption_guard.py  | 60 +++++++++++++++
 tests/ui/test_update_native.py                | 72 ++++++++++++++++--
 tests/ui_and_conv/test_prompt_tips.py         | 25 ++++---
 tests/ui_and_conv/test_update_orchestrator.py | 16 ++++
 10 files changed, 311 insertions(+), 35 deletions(-)
 create mode 100644 tests/cli/test_post_update_corruption_guard.py

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 249ba42e..3ef03fd5 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -15,6 +15,16 @@ GitHub Releases page; `0.8.0` is the new starting line.
 
 ## Unreleased
 
+- **Update notice renders below the input box.** The "Updated → vX. Restart to
+  apply." line now appears underneath the prompt's bottom border instead of inside
+  the input area.
+- **Silent native updates no longer crash a running session.** The downloaded
+  build is now staged beside the running executable and swapped in at exit (it goes
+  live on the next launch), instead of overwriting the live onefile bundle in place.
+  Overwriting it mid-session corrupted later lazy imports with a
+  `zlib.error: incorrect header check`. A boundary guard also converts any residual
+  post-update archive-corruption crash into a clear "restart to apply" message.
+
 ## 0.49.0 (2026-06-20)
 
 - **Reduction ladder in the default agent prompt.** The agent now walks an explicit,
diff --git a/src/pythinker_code/cli/__init__.py b/src/pythinker_code/cli/__init__.py
index c2625ed6..8d55da25 100644
--- a/src/pythinker_code/cli/__init__.py
+++ b/src/pythinker_code/cli/__init__.py
@@ -333,6 +333,36 @@ def _version_callback(value: bool) -> None:
         raise typer.Exit()
 
 
+def _is_post_update_bundle_corruption(exc: BaseException) -> bool:
+    """True when a frozen build's archive read failed mid-session.
+
+    This is the signature of a silent self-update having replaced the running
+    bundle on disk (see ``_install_native_archive``): the old process keeps
+    running, and the next not-yet-loaded lazy import reads from the now-stale
+    on-disk archive and raises ``zlib.error: incorrect header check``.
+
+    Gated on the PyInstaller ``sys.frozen`` marker so source / pip installs (which
+    never self-replace this way) are unaffected. The zlib error class is matched
+    by ``__module__``/``__name__`` rather than ``isinstance`` so the handler never
+    imports ``zlib`` from a possibly-corrupted archive — ``sys`` is a C built-in
+    and is always safe to import. The cause/context chain is walked with a cycle
+    guard so a self-referential chain can't spin.
+    """
+    import sys
+
+    if not getattr(sys, "frozen", False):
+        return False
+    seen: set[int] = set()
+    current: BaseException | None = exc
+    while current is not None and id(current) not in seen:
+        seen.add(id(current))
+        cls = type(current)
+        if cls.__module__ == "zlib" and cls.__name__ == "error":
+            return True
+        current = current.__cause__ or current.__context__
+    return False
+
+
 @cli.callback(invoke_without_command=True)
 def pythinker(
     ctx: typer.Context,
@@ -1310,6 +1340,17 @@ def _restore_term_and_exit(signum: int, frame: object) -> None:
             # ClickException includes the errors Typer knows how to render; don't
             # wrap them, or we'd lose the standard error UI and exit codes.
             raise
+        if _is_post_update_bundle_corruption(exc):
+            # A silent self-update replaced this running build on disk; lazy imports
+            # from the stale archive now fail. This only resolves with a restart —
+            # surface a clear, actionable message instead of a fatal traceback.
+            logger.warning("Post-update bundle corruption detected; prompting restart")
+            _emit_fatal_error(
+                "A Pythinker update was installed during this session and the "
+                "running build was replaced on disk.\n"
+                "Please restart Pythinker to finish applying the update."
+            )
+            raise typer.Exit(code=1) from exc
         logger.exception("Fatal error when running CLI")
         if debug:
             import traceback
diff --git a/src/pythinker_code/ui/shell/__init__.py b/src/pythinker_code/ui/shell/__init__.py
index 4c648235..9c6ea6ac 100644
--- a/src/pythinker_code/ui/shell/__init__.py
+++ b/src/pythinker_code/ui/shell/__init__.py
@@ -2135,7 +2135,7 @@ def _surface_installed_update_notice(self) -> None:
                 style="fg:ansiyellow",
             )
             return
-        # The persistent under-input line (_prepend_update_notice) already renders
+        # The persistent under-input line (_append_update_notice) already renders
         # the restart message; a toast duplicates it on the footer's second row.
         self._refresh_update_notice_line()
 
diff --git a/src/pythinker_code/ui/shell/prompt.py b/src/pythinker_code/ui/shell/prompt.py
index a65f4e14..37744696 100644
--- a/src/pythinker_code/ui/shell/prompt.py
+++ b/src/pythinker_code/ui/shell/prompt.py
@@ -3915,10 +3915,11 @@ def _append_history_entry(self, text: str) -> None:
                 error=exc,
             )
 
-    def _prepend_update_notice(self, fragments: list[tuple[str, str]], columns: int) -> None:
-        """Prepend a persistent yellow 'update available' line above the footer
-        separator, so it renders directly under the prompt input. No-op when no
-        update is pending; style-agnostic across both toolbar layouts."""
+    def _append_update_notice(self, fragments: list[tuple[str, str]], columns: int) -> None:
+        """Append a persistent yellow 'update available' line *below* the footer
+        separator, so it renders underneath the prompt input box rather than
+        inside it. Call this right after the separator rule. No-op when no update
+        is pending; style-agnostic across both toolbar layouts."""
         provider = getattr(self, "_update_notice_provider", None)
         if provider is None:
             return
@@ -3930,7 +3931,7 @@ def _prepend_update_notice(self, fragments: list[tuple[str, str]], columns: int)
             return
         tokens = _get_tui_tokens()
         style = f"fg:{tokens.warning or 'ansiyellow'} bold"
-        fragments[:0] = [(style, line), ("", "\n")]
+        fragments.extend([(style, line), ("", "\n")])
 
     def _render_bottom_toolbar(self) -> FormattedText:
         if (
@@ -3954,9 +3955,9 @@ def _render_bottom_toolbar(self) -> FormattedText:
         fragments: list[tuple[str, str]] = []
         tc = get_toolbar_colors()
 
-        self._prepend_update_notice(fragments, columns)
         fragments.append((self._prompt_separator_style(tc.separator), _prompt_rule(columns)))
         fragments.append(("", "\n"))
+        self._append_update_notice(fragments, columns)
 
         remaining = columns
 
@@ -4185,9 +4186,9 @@ def _render_card_bottom_toolbar(self, columns: int) -> FormattedText:
         secondary_style = f"fg:{tokens.muted}"
 
         fragments: list[tuple[str, str]] = []
-        self._prepend_update_notice(fragments, columns)
         fragments.append((self._prompt_separator_style(tc.separator), _prompt_rule(columns)))
         fragments.append(("", "\n"))
+        self._append_update_notice(fragments, columns)
 
         try:
             ctx = self._build_statusline_context(columns)
diff --git a/src/pythinker_code/ui/shell/update.py b/src/pythinker_code/ui/shell/update.py
index 98dbca31..adbd947e 100644
--- a/src/pythinker_code/ui/shell/update.py
+++ b/src/pythinker_code/ui/shell/update.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import asyncio
+import atexit
 import contextlib
 import os
 import platform
@@ -1075,6 +1076,61 @@ def _install_linux_package(asset: Path, package_kind: str) -> UpdateResult:
     return UpdateResult.UPDATED if result.returncode == 0 else UpdateResult.FAILED
 
 
+def staged_native_path() -> Path:
+    """Side path next to the running executable where a downloaded native update is
+    held until it is promoted into place. Deterministic so the orchestrator's smoke
+    check and the exit-time promotion agree on the same file."""
+    target = Path(sys.executable).resolve()
+    return target.with_name(f".{target.name}.staged")
+
+
+_staged_promotion_registered = False
+
+
+def register_staged_native_promotion() -> None:
+    """Arrange for the staged native binary to replace the running executable at
+    process exit.
+
+    Swapping at exit — rather than in place mid-session — is the whole point: the
+    running onefile build reads its Python archive lazily from ``sys.executable``,
+    so overwriting that path while the process is alive corrupts later imports
+    (``zlib.error: incorrect header check``). At exit no further imports happen, so
+    the swap is safe, and the new binary goes live on the next launch — exactly what
+    the "Updated → vX. Restart to apply." notice promises. Registration is idempotent
+    so a silent update followed by a manual ``/update`` only swaps once.
+    """
+    global _staged_promotion_registered
+    if _staged_promotion_registered:
+        return
+    atexit.register(_promote_staged_native_update)
+    _staged_promotion_registered = True
+
+
+def discard_staged_native_update() -> None:
+    """Remove a staged binary that must not be promoted (e.g. it failed its smoke
+    check). Fail-open: a leftover staged file is only ever promoted deliberately."""
+    with contextlib.suppress(OSError):
+        staged_native_path().unlink()
+
+
+def _promote_staged_native_update() -> None:
+    # ponytail: known residual windows, both far smaller than the mid-session login
+    # crash this replaces — (1) a hard kill (SIGKILL) skips atexit, so promotion
+    # defers to the next clean exit (self-healing); (2) a first-time lazy import
+    # during the rest of interpreter shutdown, after the swap below, could read
+    # stale bytes. Closing both needs a detached post-exit swapper (the deferred
+    # seamless-relaunch design); not built here.
+    staged = staged_native_path()
+    if not staged.is_file():
+        return
+    target = Path(sys.executable).resolve()
+    try:
+        os.replace(staged, target)
+    except OSError:
+        # Leave the staged binary in place; a later clean exit retries the swap.
+        logger.exception("Failed to promote staged native update on exit:")
+
+
 def _install_native_archive(archive: Path) -> UpdateResult:
     target = Path(sys.executable).resolve()
     extract_dir = archive.parent / "extract"
@@ -1091,18 +1147,23 @@ def _install_native_archive(archive: Path) -> UpdateResult:
         logger.error("Native archive did not contain a pythinker executable")
         return UpdateResult.FAILED
 
-    replacement = target.with_name(f".{target.name}.new-{os.getpid()}")
+    # Stage the new binary beside the running executable instead of overwriting it
+    # in place. Promotion into `target` happens at process exit (see
+    # register_staged_native_promotion), keeping the live build's on-disk archive
+    # intact so mid-session lazy imports never read stale bytes.
+    staged = staged_native_path()
+    staging_tmp = target.with_name(f".{target.name}.new-{os.getpid()}")
     try:
-        shutil.copyfile(extracted, replacement)
-        replacement.chmod(target.stat().st_mode | 0o755)
-        os.replace(replacement, target)
+        shutil.copyfile(extracted, staging_tmp)
+        staging_tmp.chmod(target.stat().st_mode | 0o755)
+        os.replace(staging_tmp, staged)
         (target.parent / ".pythinker-native").write_text(
             "pythinker-native-build\n", encoding="utf-8"
         )
     except OSError:
-        logger.exception("Failed to replace native executable:")
+        logger.exception("Failed to stage native executable:")
         with contextlib.suppress(OSError):
-            replacement.unlink()
+            staging_tmp.unlink()
         return UpdateResult.FAILED
     return UpdateResult.UPDATED
 
diff --git a/src/pythinker_code/ui/shell/update_orchestrator.py b/src/pythinker_code/ui/shell/update_orchestrator.py
index d4ed009e..98236e57 100644
--- a/src/pythinker_code/ui/shell/update_orchestrator.py
+++ b/src/pythinker_code/ui/shell/update_orchestrator.py
@@ -382,8 +382,11 @@ async def run_update_job(
             if smoke_ok:
                 message = smoke_message
                 _write_last_success(job_id=job_id, message=message)
+                _finalize_native_staging(promote=True)
             else:
                 message = f"{SMOKE_CHECK_FAILED_PREFIX}{smoke_message}"
+                # Never promote a staged binary that can't even print --version.
+                _finalize_native_staging(promote=False)
 
         write_update_status(
             _new_status(
@@ -432,10 +435,37 @@ def _write_last_success(*, job_id: str, message: str) -> None:
 
 def _smoke_check_command() -> list[str]:
     if is_native_build():
-        return [sys.executable, "--version"]
+        # A native update is staged beside the running exe, not swapped in place,
+        # so validate the staged binary — checking the still-running old exe would
+        # prove nothing about the update.
+        from pythinker_code.ui.shell.update import staged_native_path
+
+        staged = staged_native_path()
+        exe = str(staged) if staged.is_file() else sys.executable
+        return [exe, "--version"]
     return [sys.executable, "-P", "-m", "pythinker_code", "--version"]
 
 
+def _finalize_native_staging(*, promote: bool) -> None:
+    """After the smoke check, either arm the staged native binary for exit-time
+    promotion (it ran) or discard it (it failed). No-op for non-native installs and
+    when nothing was staged."""
+    if not is_native_build():
+        return
+    from pythinker_code.ui.shell.update import (
+        discard_staged_native_update,
+        register_staged_native_promotion,
+        staged_native_path,
+    )
+
+    if not staged_native_path().is_file():
+        return
+    if promote:
+        register_staged_native_promotion()
+    else:
+        discard_staged_native_update()
+
+
 def _smoke_check_cwd() -> Path:
     try:
         return Path(sys.executable).resolve().parent
diff --git a/tests/cli/test_post_update_corruption_guard.py b/tests/cli/test_post_update_corruption_guard.py
new file mode 100644
index 00000000..ac29d540
--- /dev/null
+++ b/tests/cli/test_post_update_corruption_guard.py
@@ -0,0 +1,60 @@
+"""Guard that turns post-self-update archive corruption into a restart prompt.
+
+A silent native update replaces the running bundle on disk; the next not-yet-loaded
+lazy import then reads a stale archive and raises ``zlib.error``. The CLI boundary
+must recognize that signature (frozen build + zlib error in the cause chain) and
+surface a clear restart message instead of a fatal traceback — without re-importing
+anything from the corrupted archive.
+"""
+
+import zlib
+
+import pytest
+
+from pythinker_code.cli import _is_post_update_bundle_corruption
+
+
+@pytest.fixture
+def frozen(monkeypatch):
+    monkeypatch.setattr("sys.frozen", True, raising=False)
+
+
+def test_detects_direct_zlib_error_on_frozen_build(frozen):
+    assert _is_post_update_bundle_corruption(zlib.error("incorrect header check")) is True
+
+
+def test_detects_zlib_error_in_cause_chain(frozen):
+    try:
+        try:
+            raise zlib.error("incorrect header check")
+        except zlib.error as inner:
+            raise ImportError("cannot load module") from inner
+    except ImportError as exc:
+        assert _is_post_update_bundle_corruption(exc) is True
+
+
+def test_detects_zlib_error_in_implicit_context(frozen):
+    try:
+        try:
+            raise zlib.error("incorrect header check")
+        except zlib.error:
+            raise RuntimeError("boom")  # noqa: B904 - implicit __context__ is the point
+    except RuntimeError as exc:
+        assert _is_post_update_bundle_corruption(exc) is True
+
+
+def test_ignores_zlib_error_when_not_frozen(monkeypatch):
+    monkeypatch.delattr("sys.frozen", raising=False)
+    assert _is_post_update_bundle_corruption(zlib.error("incorrect header check")) is False
+
+
+def test_ignores_unrelated_error_on_frozen_build(frozen):
+    assert _is_post_update_bundle_corruption(ValueError("nope")) is False
+
+
+def test_handles_cyclic_cause_chain(frozen):
+    a = RuntimeError("a")
+    b = RuntimeError("b")
+    a.__cause__ = b
+    b.__cause__ = a  # cycle: must terminate, not spin
+    assert _is_post_update_bundle_corruption(a) is False
diff --git a/tests/ui/test_update_native.py b/tests/ui/test_update_native.py
index 293c1986..a8c5a618 100644
--- a/tests/ui/test_update_native.py
+++ b/tests/ui/test_update_native.py
@@ -67,23 +67,79 @@ def test_native_prompt_does_not_leak_marker(monkeypatch):
     assert "native updater" in text
 
 
-def test_install_native_archive_replaces_current_executable(monkeypatch, tmp_path):
-    current = tmp_path / "pythinker"
-    current.write_text("old", encoding="utf-8")
-    current.chmod(0o755)
-
+def _make_native_archive(tmp_path, body: str = "new"):
     payload = tmp_path / "payload"
     payload.mkdir()
-    (payload / "pythinker").write_text("new", encoding="utf-8")
+    (payload / "pythinker").write_text(body, encoding="utf-8")
     archive = tmp_path / "pythinker-0.2.0-x86_64-unknown-linux-gnu.tar.gz"
     with tarfile.open(archive, "w:gz") as tar:
         tar.add(payload / "pythinker", arcname="pythinker")
+    return archive
+
+
+def test_install_native_archive_stages_without_touching_running_exe(monkeypatch, tmp_path):
+    # The running executable must be left untouched (overwriting it mid-session
+    # corrupts this onefile build's lazy archive reads). The new binary lands in
+    # the staged side path instead, ready to promote on exit.
+    current = tmp_path / "pythinker"
+    current.write_text("old", encoding="utf-8")
+    current.chmod(0o755)
+    archive = _make_native_archive(tmp_path)
 
     monkeypatch.setattr(upd.sys, "executable", str(current))
 
     assert upd._install_native_archive(archive) is upd.UpdateResult.UPDATED
-    assert current.read_text(encoding="utf-8") == "new"
-    assert current.stat().st_mode & 0o111
+    # Running exe is unchanged...
+    assert current.read_text(encoding="utf-8") == "old"
+    # ...the update is staged beside it, executable, and marked native.
+    staged = upd.staged_native_path()
+    assert staged.read_text(encoding="utf-8") == "new"
+    assert staged.stat().st_mode & 0o111
     assert (tmp_path / ".pythinker-native").read_text(encoding="utf-8") == (
         "pythinker-native-build\n"
     )
+
+
+def test_promote_staged_native_update_swaps_into_place(monkeypatch, tmp_path):
+    current = tmp_path / "pythinker"
+    current.write_text("old", encoding="utf-8")
+    monkeypatch.setattr(upd.sys, "executable", str(current))
+    upd.staged_native_path().write_text("new", encoding="utf-8")
+
+    upd._promote_staged_native_update()
+
+    assert current.read_text(encoding="utf-8") == "new"
+    assert not upd.staged_native_path().exists()
+
+
+def test_promote_staged_native_update_noop_without_staged_file(monkeypatch, tmp_path):
+    current = tmp_path / "pythinker"
+    current.write_text("old", encoding="utf-8")
+    monkeypatch.setattr(upd.sys, "executable", str(current))
+
+    upd._promote_staged_native_update()  # no staged file → no-op
+
+    assert current.read_text(encoding="utf-8") == "old"
+
+
+def test_discard_staged_native_update_removes_file(monkeypatch, tmp_path):
+    current = tmp_path / "pythinker"
+    current.write_text("old", encoding="utf-8")
+    monkeypatch.setattr(upd.sys, "executable", str(current))
+    upd.staged_native_path().write_text("broken", encoding="utf-8")
+
+    upd.discard_staged_native_update()
+
+    assert not upd.staged_native_path().exists()
+    assert current.read_text(encoding="utf-8") == "old"  # running exe untouched
+
+
+def test_register_staged_native_promotion_is_idempotent(monkeypatch):
+    registered: list[object] = []
+    monkeypatch.setattr(upd.atexit, "register", lambda fn: registered.append(fn))
+    monkeypatch.setattr(upd, "_staged_promotion_registered", False)
+
+    upd.register_staged_native_promotion()
+    upd.register_staged_native_promotion()
+
+    assert registered == [upd._promote_staged_native_update]
diff --git a/tests/ui_and_conv/test_prompt_tips.py b/tests/ui_and_conv/test_prompt_tips.py
index 6666966a..6640cfd7 100644
--- a/tests/ui_and_conv/test_prompt_tips.py
+++ b/tests/ui_and_conv/test_prompt_tips.py
@@ -1711,28 +1711,29 @@ async def prompt_async(self, **kwargs):
     assert prompt_session.last_submission_was_running is False
 
 
-def test_prepend_update_notice_inserts_line_above_separator():
+def test_append_update_notice_inserts_line_below_separator():
+    # Caller appends the separator rule + newline first, then the notice — so the
+    # notice renders underneath the input box (below the bottom border), not inside it.
     fragments: list[tuple[str, str]] = [("sep", "────"), ("", "\n")]
     fake = SimpleNamespace(_update_notice_provider=lambda: "↑ Update available — v9.9.9 · /update")
-    CustomPromptSession._prepend_update_notice(cast(Any, fake), fragments, 80)
-    # Notice is the first row (its own line), then a newline, then the original
-    # separator — i.e. it renders directly under the input, above the footer rule.
-    assert "Update available" in fragments[0][1]
-    assert "v9.9.9" in fragments[0][1]
-    assert "bold" in fragments[0][0]
+    CustomPromptSession._append_update_notice(cast(Any, fake), fragments, 80)
+    assert fragments[0] == ("sep", "────")
     assert fragments[1] == ("", "\n")
-    assert fragments[2] == ("sep", "────")
+    assert "Update available" in fragments[2][1]
+    assert "v9.9.9" in fragments[2][1]
+    assert "bold" in fragments[2][0]
+    assert fragments[3] == ("", "\n")
 
 
-def test_prepend_update_notice_noop_when_no_update():
+def test_append_update_notice_noop_when_no_update():
     fragments: list[tuple[str, str]] = [("sep", "────")]
     fake = SimpleNamespace(_update_notice_provider=lambda: None)
-    CustomPromptSession._prepend_update_notice(cast(Any, fake), fragments, 80)
+    CustomPromptSession._append_update_notice(cast(Any, fake), fragments, 80)
     assert fragments == [("sep", "────")]
 
 
-def test_prepend_update_notice_noop_when_no_provider():
+def test_append_update_notice_noop_when_no_provider():
     fragments: list[tuple[str, str]] = [("sep", "x")]
     fake = SimpleNamespace(_update_notice_provider=None)
-    CustomPromptSession._prepend_update_notice(cast(Any, fake), fragments, 80)
+    CustomPromptSession._append_update_notice(cast(Any, fake), fragments, 80)
     assert fragments == [("sep", "x")]
diff --git a/tests/ui_and_conv/test_update_orchestrator.py b/tests/ui_and_conv/test_update_orchestrator.py
index 5ea3da7f..88a24ab9 100644
--- a/tests/ui_and_conv/test_update_orchestrator.py
+++ b/tests/ui_and_conv/test_update_orchestrator.py
@@ -272,6 +272,22 @@ def test_native_smoke_check_does_not_use_python_module_import(monkeypatch):
     assert orchestrator._smoke_check_command() == ["/opt/pythinker/pythinker", "--version"]
 
 
+def test_native_smoke_check_targets_staged_binary_when_present(monkeypatch, tmp_path):
+    # A native update is staged beside the running exe; the smoke check must
+    # validate the staged binary, not the still-running old one.
+    current = tmp_path / "pythinker"
+    current.write_text("old", encoding="utf-8")
+    monkeypatch.setattr(orchestrator, "is_native_build", lambda: True)
+    monkeypatch.setattr(orchestrator.sys, "executable", str(current))
+
+    from pythinker_code.ui.shell.update import staged_native_path
+
+    staged = staged_native_path()
+    staged.write_text("new", encoding="utf-8")
+
+    assert orchestrator._smoke_check_command() == [str(staged), "--version"]
+
+
 @pytest.mark.asyncio
 async def test_do_update_mirrors_messages_to_output_callback(monkeypatch, tmp_path):
     messages: list[str] = []

From 1714120d58304c32c845205dd8906fbbd1484c2b Mon Sep 17 00:00:00 2001
From: Mohamed Elkholy <moelkholy1995@gmail.com>
Date: Sat, 20 Jun 2026 13:14:31 -0400
Subject: [PATCH 2/2] =?UTF-8?q?fix(update):=20address=20review=20=E2=80=94?=
 =?UTF-8?q?=20narrow=20corruption=20classifier,=20drop=20dedup=20flag?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- _is_post_update_bundle_corruption now also requires the documented
  "incorrect header check" message, so an unrelated zlib decompression failure
  on a frozen build isn't misclassified as a stale bundle and masked behind a
  restart-only message (CodeRabbit).
- Drop the _staged_promotion_registered module global; register_staged_native
  _promotion now registers unconditionally and relies on the already-idempotent
  atexit handler (no-ops once the staged file is promoted), removing the unused-
  global finding (github-code-quality).
- Tests updated/added: unrelated-zlib-message rejection; handler idempotency on
  double-run; registration registers the atexit handler.
---
 src/pythinker_code/cli/__init__.py            | 19 +++++++++++++------
 src/pythinker_code/ui/shell/update.py         | 14 +++++---------
 .../cli/test_post_update_corruption_guard.py  |  6 ++++++
 tests/ui/test_update_native.py                | 19 ++++++++++++++++---
 4 files changed, 40 insertions(+), 18 deletions(-)

diff --git a/src/pythinker_code/cli/__init__.py b/src/pythinker_code/cli/__init__.py
index 8d55da25..1c81a620 100644
--- a/src/pythinker_code/cli/__init__.py
+++ b/src/pythinker_code/cli/__init__.py
@@ -342,11 +342,14 @@ def _is_post_update_bundle_corruption(exc: BaseException) -> bool:
     on-disk archive and raises ``zlib.error: incorrect header check``.
 
     Gated on the PyInstaller ``sys.frozen`` marker so source / pip installs (which
-    never self-replace this way) are unaffected. The zlib error class is matched
-    by ``__module__``/``__name__`` rather than ``isinstance`` so the handler never
-    imports ``zlib`` from a possibly-corrupted archive — ``sys`` is a C built-in
-    and is always safe to import. The cause/context chain is walked with a cycle
-    guard so a self-referential chain can't spin.
+    never self-replace this way) are unaffected, and narrowed to the documented
+    corruption message (``incorrect header check``) so an unrelated decompression
+    failure elsewhere on a frozen build is *not* misclassified as a stale bundle and
+    masked behind a restart-only message. The zlib error class is matched by
+    ``__module__``/``__name__`` rather than ``isinstance`` so the handler never
+    imports ``zlib`` from a possibly-corrupted archive — ``sys`` is a C built-in and
+    is always safe to import. The cause/context chain is walked with a cycle guard so
+    a self-referential chain can't spin.
     """
     import sys
 
@@ -357,7 +360,11 @@ def _is_post_update_bundle_corruption(exc: BaseException) -> bool:
     while current is not None and id(current) not in seen:
         seen.add(id(current))
         cls = type(current)
-        if cls.__module__ == "zlib" and cls.__name__ == "error":
+        if (
+            cls.__module__ == "zlib"
+            and cls.__name__ == "error"
+            and "incorrect header check" in str(current).lower()
+        ):
             return True
         current = current.__cause__ or current.__context__
     return False
diff --git a/src/pythinker_code/ui/shell/update.py b/src/pythinker_code/ui/shell/update.py
index adbd947e..8e83c69d 100644
--- a/src/pythinker_code/ui/shell/update.py
+++ b/src/pythinker_code/ui/shell/update.py
@@ -1084,9 +1084,6 @@ def staged_native_path() -> Path:
     return target.with_name(f".{target.name}.staged")
 
 
-_staged_promotion_registered = False
-
-
 def register_staged_native_promotion() -> None:
     """Arrange for the staged native binary to replace the running executable at
     process exit.
@@ -1096,14 +1093,13 @@ def register_staged_native_promotion() -> None:
     so overwriting that path while the process is alive corrupts later imports
     (``zlib.error: incorrect header check``). At exit no further imports happen, so
     the swap is safe, and the new binary goes live on the next launch — exactly what
-    the "Updated → vX. Restart to apply." notice promises. Registration is idempotent
-    so a silent update followed by a manual ``/update`` only swaps once.
+    the "Updated → vX. Restart to apply." notice promises.
+
+    Registering more than once (e.g. a silent update plus a manual ``/update`` in the
+    same session) is harmless: the handler no-ops once the staged file has been
+    promoted, so any duplicate registration just runs a second no-op.
     """
-    global _staged_promotion_registered
-    if _staged_promotion_registered:
-        return
     atexit.register(_promote_staged_native_update)
-    _staged_promotion_registered = True
 
 
 def discard_staged_native_update() -> None:
diff --git a/tests/cli/test_post_update_corruption_guard.py b/tests/cli/test_post_update_corruption_guard.py
index ac29d540..b5f2048e 100644
--- a/tests/cli/test_post_update_corruption_guard.py
+++ b/tests/cli/test_post_update_corruption_guard.py
@@ -52,6 +52,12 @@ def test_ignores_unrelated_error_on_frozen_build(frozen):
     assert _is_post_update_bundle_corruption(ValueError("nope")) is False
 
 
+def test_ignores_zlib_error_with_unrelated_message(frozen):
+    # A decompression failure that is NOT the bundle-corruption signature (e.g. a
+    # bad gzip HTTP response) must not be masked behind a restart-only message.
+    assert _is_post_update_bundle_corruption(zlib.error("invalid distance too far back")) is False
+
+
 def test_handles_cyclic_cause_chain(frozen):
     a = RuntimeError("a")
     b = RuntimeError("b")
diff --git a/tests/ui/test_update_native.py b/tests/ui/test_update_native.py
index a8c5a618..11c81d65 100644
--- a/tests/ui/test_update_native.py
+++ b/tests/ui/test_update_native.py
@@ -134,12 +134,25 @@ def test_discard_staged_native_update_removes_file(monkeypatch, tmp_path):
     assert current.read_text(encoding="utf-8") == "old"  # running exe untouched
 
 
-def test_register_staged_native_promotion_is_idempotent(monkeypatch):
+def test_register_staged_native_promotion_registers_atexit_handler(monkeypatch):
     registered: list[object] = []
     monkeypatch.setattr(upd.atexit, "register", lambda fn: registered.append(fn))
-    monkeypatch.setattr(upd, "_staged_promotion_registered", False)
 
-    upd.register_staged_native_promotion()
     upd.register_staged_native_promotion()
 
     assert registered == [upd._promote_staged_native_update]
+
+
+def test_promote_staged_native_update_is_idempotent_when_run_twice(monkeypatch, tmp_path):
+    # Duplicate registration is safe because the handler itself is idempotent: the
+    # second run finds the staged file already promoted and no-ops without error.
+    current = tmp_path / "pythinker"
+    current.write_text("old", encoding="utf-8")
+    monkeypatch.setattr(upd.sys, "executable", str(current))
+    upd.staged_native_path().write_text("new", encoding="utf-8")
+
+    upd._promote_staged_native_update()
+    upd._promote_staged_native_update()
+
+    assert current.read_text(encoding="utf-8") == "new"
+    assert not upd.staged_native_path().exists()