Skip to content

perf(theme): defer heavy imports until user actions need them#668

Open
thomwebb wants to merge 7 commits into
mpfaffenberger:mainfrom
thomwebb:perf/theme-lazy-imports
Open

perf(theme): defer heavy imports until user actions need them#668
thomwebb wants to merge 7 commits into
mpfaffenberger:mainfrom
thomwebb:perf/theme-lazy-imports

Conversation

@thomwebb

@thomwebb thomwebb commented Jul 23, 2026

Copy link
Copy Markdown
Contributor

Summary

Move eight sibling imports in code_puppy/plugins/theme/register_callbacks.py
from module top-level into the functions that actually use them, and
replace one eagerly-registered callback with a lazy shim.

The theme plugin's startup callback only needs to apply a default theme
on first run. Everything else — .picker, .themes,
.prompt_toolkit_theme, .bundled_palettes, plus a handful of heavier
code_puppy internals — is only needed when the user runs /theme or
triggers a render.

Measured impact

Honest numbers — this is a modest win:

Metric Before After
register_callbacks isolated import (warm) ~21 ms ~11 ms
register_callbacks isolated import (cold) ~82 ms ~53 ms
End-to-end --version (warm) 1.78 s 1.75 s

The end-to-end saving is within measurement noise (~30 ms). Verified via a
sys.modules snapshot after plugins.load_plugin_callbacks() completes:

code_puppy.plugins.theme.themes:               not loaded (good)
code_puppy.plugins.theme.picker:               not loaded (good)
code_puppy.plugins.theme.prompt_toolkit_theme: not loaded (good)
code_puppy.plugins.theme.bundled_palettes:    not loaded (good)

Scope caveat (honest correction)

An earlier draft of this PR claimed the theme plugin "no longer
transitively pulls in prompt_toolkit at plugin-discovery time." That
claim was false and has been retracted from the in-file comment and
this description.

Repro:

from code_puppy.callbacks import register_callback  # baseline
before = set(sys.modules)
import code_puppy.plugins.theme  # just the package
# → 115 prompt_toolkit.* modules already loaded

code_puppy/plugins/theme/__init__.py eagerly imports .content_styles,
.osc_palette, and .rich_themes, and calls reapply_from_config() on
each at package-import time. That path transitively pulls in
code_puppy.messaging.rich_renderer, which loads ~115 prompt_toolkit
modules. Nothing this PR does to register_callbacks.py can shortcut
that — those modules are already resident by the time register_callbacks
is imported.

The measured 10-30 ms savings this PR delivers therefore come from
deferring .picker (TUI widgets), .themes (large color-menu tables),
and .prompt_toolkit_theme — not from avoiding prompt_toolkit
altogether. A future PR could defer the __init__.py reapply sweep to
a startup callback, which would recover the prompt_toolkit cost for
--version / --help invocations that skip startup. Out of scope here.

Why we did it anyway

  • Establishes a clean lazy-import pattern for other plugins to follow.
  • The measured 10-30 ms is real, if modest.
  • Test suite got stronger: 9 white-box patches previously targeted
    register_callbacks.X (fragile — patching a re-import
    point instead of the source). They now target the actual source
    modules, which is the healthier shape.
  • Adds a _prompt_toolkit_style lazy shim with a strict signature
    ((style: BaseStyle | None) -> BaseStyle) so
    merge_with_active_style gets called on-demand rather than requiring
    its module at registration time.

Tests

  • All 110 theme-related tests pass (test_theme_plugin.py +
    test_prompt_toolkit_theme.py).
  • Retargeted 9 white-box patches from register_callbacks.X (which no
    longer exists at module scope after this change) to the actual source
    modules (code_puppy.config.get_value,
    code_puppy.plugins.theme.osc_palette.get_saved_palette, etc.).
  • Added two new tests covering the _prompt_toolkit_style shim: one
    verifies it delegates its single argument to
    merge_with_active_style (regression guard against silent signature
    widening), and one confirms the shim preserves __name__ for error
    logging.
  • ruff check and ruff format --check clean.

Broader startup-perf context

While profiling this I measured the import cost of each plugin's
register_callbacks in isolation (snapshot sys.modules before/after,
bucket the delta by heavy library). A few results worth passing along:

Plugin Modules added Notable heavies
claude_code_oauth +2173 openai(557) anthropic(445) pydantic_ai(56) playwright(53)
theme (before this PR) +940 pydantic_ai(52) prompt_toolkit(115) rich(60)
aws_bedrock +932 pydantic_ai(52) playwright(53)
obsidian_agent +378 pydantic_ai(19) mcp(82)

Takeaways for future PRs (not part of this change):

  • pydantic_ai (~52 modules) loads per-plugin. A single root import
    point is the leverage lever — likely via the type imports referenced
    at module scope in the agent-registration API.
  • Two plugins unexpectedly pull in playwright (aws_bedrock,
    claude_code_oauth). Neither does browser automation. Likely a
    code_puppy.tools-registry-level bug where registering browser tools
    requires importing playwright.
  • claude_code_oauth pulls in the entire OpenAI SDK (557 modules).
    A plugin for Claude authentication should not need OpenAI's Python
    client.
  • The dominant wall-clock cost at startup is the synchronous pypi.org
    version check
    — ~0.8 s median on a healthy network, up to
    10-plus seconds when pypi.org is slow or blocked by a corporate
    proxy. Kicking it to a background task is the single biggest UX lever
    and doesn't require touching plugin structure.

Happy to file follow-up issues on this repo for any of the above if
they'd be useful — I've been tracking them internally.

Files touched

  • code_puppy/plugins/theme/register_callbacks.py — 8 sibling imports
    moved into functions; _prompt_toolkit_style lazy shim added; NOTE
    comment rewritten for accuracy.
  • tests/plugins/test_theme_plugin.py — patch targets retargeted from
    register_callbacks.X to source modules; two new tests for the
    shim added.

TJ Webb added 4 commits July 23, 2026 09:39
The theme plugin's register_callbacks.py eagerly imported eight sibling
modules at the top of the file — including .picker (prompt_toolkit
widgets), .themes (728 lines), and .prompt_toolkit_theme. None were
needed to satisfy the startup callback, whose only job is to apply the
default theme on a fresh install.

Move all sibling imports into the functions that actually use them.
The prompt_toolkit_style callback needed a small shim because the
callback registry captures the callable at registration time.

Verified via sys.modules snapshot: themes.py, picker.py,
prompt_toolkit_theme.py, and bundled_palettes.py no longer appear in
sys.modules after plugins.load_plugin_callbacks() completes.

Wall-clock impact is modest:
  - Isolated register_callbacks import: ~30ms cold, ~10ms warm
    (down from ~82ms cold / ~21ms warm)
  - End-to-end --version: within measurement noise (~30ms)

The bigger win is architectural: the theme package no longer
transitively requires prompt_toolkit at plugin-discovery time, and
this establishes the lazy-import pattern for other plugins to follow.

Tests: retargeted white-box patches from register_callbacks.X to the
actual source modules (config.get_value, osc_palette.get_saved_palette,
themes.apply, etc.). All 121 theme-related tests pass.

Report: docs/STARTUP_PERFORMANCE.md documents the full analysis and
remaining priorities (P2 — skip plugin discovery on --version/--help —
is now clearly the biggest lever).
…h profiling

The original priorities were scored against `--version` timing, which
turned out to be a bad proxy for what humans actually experience.
Interactive-launch profiling shows:

- pypi.org version check is the biggest lever (0.8s median, up to 13-18s
  when pypi is unreachable through the Walmart sysproxy). Elevated from
  P5 to P1.

- The 'skip plugins on --version' fix (was P2) does NOT help interactive
  users at all. Downgraded to P4.

- Per-plugin isolated-import profiling surfaced two bugs:
  * claude_code_oauth pulls in the entire OpenAI SDK (557 modules)
  * Multiple plugins (aws_bedrock, claude_code_oauth) transitively
    load playwright via the tools registry

- pydantic_ai (~52 modules per plugin) is the highest-leverage single
  target for lazy provider-SDK loading.

Adds:
- Revised priority table at the top of the report, superseding the old
  P1-P5 ordering below.
- Per-plugin isolated-import profile with heavy-library breakdown.
- Documented methodology for interactive-launch and per-plugin
  measurement.

The old sections are kept for context but marked as superseded.
CI caught one `with patch(...) as ...` block that ruff format wanted
collapsed onto a single line. Local `ruff format .` before the initial
commit missed this file because the format ran before the patch
retargets landed.
The perf-plan doc is being moved out of this PR. The theme lazy-import
code change (ccd59cb) is the only substantive change; key findings and
follow-up context are now inlined in the PR description instead.
@thomwebb thomwebb changed the title perf(theme): defer heavy imports + revise startup perf analysis perf(theme): defer heavy imports until user actions need them Jul 23, 2026
TJ Webb added 3 commits July 23, 2026 10:29
The note referenced docs/STARTUP_PERFORMANCE.md (dropped from this PR)
and quoted the ~267ms figure that was later corrected as measurement
noise. Replace with a self-contained explanation that describes what
the pattern does and why (avoid pulling in prompt_toolkit at plugin
discovery + short-circuit on already-configured installs).
These two user-facing strings had gotten converted to a \U0001f3a8
Unicode escape sequence during the earlier refactor pass. Semantically
identical at runtime, but gratuitously different from the upstream
source style. Restore the literal character to keep this PR's diff
focused on the actual import motions.
Fixes uncovered by an adversarial review pass:

* Correct the module NOTE comment. The prior version claimed the theme
  plugin "no longer transitively requires prompt_toolkit at plugin
  discovery," which is factually false — theme/__init__.py eagerly
  imports .content_styles / .osc_palette / .rich_themes and calls
  reapply_from_config() on each, which drags in ~115 prompt_toolkit
  modules via code_puppy.messaging.rich_renderer. The rewrite scopes
  the win honestly (.picker + .themes + .prompt_toolkit_theme stay off
  the discovery path; "never loaded on the common launch path" softened
  to reflect that REPL renders still trigger them on first paint) and
  flags the __init__.py deferral as a follow-up.

* Tighten _prompt_toolkit_style shim signature. The previous
  '*args, **kwargs' shape silently swallowed bogus calls that would
  have TypeError'd against the real callable. Replaced with an explicit
  '(style: BaseStyle | None = None) -> BaseStyle' signature using
  TYPE_CHECKING so no runtime import cost is added. Also preserves the
  real symbol name via __name__ assignment so
  code_puppy.callbacks failure logs remain useful.

* Push imports in _apply_default_theme_on_first_run below the
  early-return check so the common case (theme already configured)
  doesn't bind unnecessary symbols. Consistent with the PR's philosophy.

* Add two new unit tests: one verifying the shim delegates its single
  argument to merge_with_active_style with no silent widening, and one
  confirming __name__ still reports the real callback symbol.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant