Overlap embedding warm-up with chat restoration

Load sentence-transformers on the existing cache miss and construct the
preloaded model off the shared background event loop. Start warm-up before
chat restoration inside extensible init_a0, preserving start-hook control
and sharing the loaded model with ordinary runtime requests.

Remove unused speech preload imports and dead warm-up functions. Seed APT
indexes only when absent, skipping the refresh on normal container boots.
Keep the existing Python tunnel service and its provider lifecycle.

Validation: 49 focused framework tests passed, including extension-hook
short-circuit behavior, concurrent embedding-cache reuse, APT cache states,
tunnel behavior and CSRF origins. Three cold Fable chat trials per version:
median submit-to-first-token 9.588s -> 4.234s;
local preparation 6.904s -> 1.197s;
HTTP readiness 42.637s -> 41.139s;
boot-to-first-token 52.093s -> 46.657s.
Provider timing varies; these are local repeated measurements. The live
baseline already omitted APT refresh, so it receives no timing credit.

AGENTS.md

@@ -22,6 +22,7 @@
 - Native function calls use canonical call content for history and combine public response text with calls for repeat comparison. The canonical history content passes through the normal history template/masking hook; provider output metadata remains intact.
 - `initialize.py` owns framework initialization.
 - `models.py` owns model-provider configuration and LiteLLM integration. `get_api_key_raw()` reads the stored key or key list for settings/editing; `get_api_key()` selects a runtime key and runs provider extension hooks.
+- Local sentence-transformers load only on an embedding cache miss, not when importing `models`. `run_ui.init_a0` starts the deferred preload before chat restoration, overlapping model loading while preserving the initialization extension hooks. `preload.py` constructs the configured local embedding model in a worker thread; it must not import unused speech runtimes or block the shared background event loop.
 - `models._merge_provider_defaults` merges YAML-declared `extra_body` defaults with legacy top-level values for those declared fields and explicit nested overrides, in that order; unrelated kwargs and original-provider API-key selection remain intact.
 - `models.apply_rate_limiter` skips token counting and returns no limiter when all request, input, and output limits are unset; any configured limit retains normal accounting.
 - `requirements.txt` includes HTTPX's SOCKS extra so framework clients support SOCKS proxy environment variables before plugin setup.

docker/run/AGENTS.md

@@ -22,6 +22,7 @@
 - For v2.13 builds, retain the Python 3.13-compatible ATK, LibreOffice/UNO, and Xpra pins in `fs/ins/install_additional.sh`. Resolve them from the signed Debian `forky` archive dated `20260624T000000Z` and Xpra `trixie`, using build-scoped APT options that preserve normal Kali sources.
 - Do not bake secrets, local `.env` values, or user data into the image.
 - Runtime startup must ensure `/a0/usr/uploads` exists before supervised services start.
+- Runtime startup seeds APT package lists in the background only when no package index is cached (fresh images clear these at build time). Normal restarts skip the refresh; installation scripts and plugin hooks refresh metadata when installing packages.
 - Runtime startup raises the soft open-file limit toward `A0_NOFILE_LIMIT` (default `65535`) before supervisord starts, bounded by the container hard limit.
 - Self-update user-data backups skip Time Travel shadow history under `usr/.time_travel/` and transient Desktop agent state.
 - Self-update rollback stashes use immutable Git object IDs for creation checks and restoration; resolve the matching reflog selector only when dropping that stash, preserving unrelated entries. Failed restoration retains the stash.

docker/run/fs/exe/initialize.sh

@@ -44,6 +44,12 @@ raise_open_file_limit() {
     fi
 }
 
+initialize_package_lists() {
+    if ! compgen -G "/var/lib/apt/lists/*_Packages*" > /dev/null; then
+        apt-get update > /dev/null 2>&1 &
+    fi
+}
+
 raise_open_file_limit
 
 # Copy all contents from persistent /per to root directory (/) without overwriting
@@ -56,8 +62,7 @@ mkdir -p /a0/usr/uploads
 chmod 444 /root/.bashrc
 chmod 444 /root/.profile
 
-# update package list to save time later
-apt-get update > /dev/null 2>&1 &
+initialize_package_lists
 
 # let supervisord handle the services
 exec /usr/bin/supervisord -c /etc/supervisor/conf.d/supervisord.conf

models.py

@@ -14,6 +14,7 @@ from typing import (
     AsyncIterator,
     Tuple,
     TypedDict,
+    TYPE_CHECKING,
 )
 
 from litellm import embedding
@@ -43,9 +44,11 @@ from langchain_core.messages import (
     SystemMessage,
 )
 from langchain.embeddings.base import Embeddings
-from sentence_transformers import SentenceTransformer
 from pydantic import ConfigDict
 
+if TYPE_CHECKING:
+    from sentence_transformers import SentenceTransformer
+
 
 DEFAULT_LITELLM_GLOBAL_KWARGS: dict[str, Any] = {
     "drop_params": True,
@@ -832,17 +835,19 @@ class LiteLLMEmbeddingWrapper(Embeddings):
         return self.embed([text])[0]
 
 
-_LOCAL_EMBEDDING_MODELS: dict[tuple[str, str], SentenceTransformer] = {}
+_LOCAL_EMBEDDING_MODELS: dict[tuple[str, str], "SentenceTransformer"] = {}
 _LOCAL_EMBEDDING_MODELS_LOCK = threading.Lock()
 
 
 def _get_local_embedding_model(
     model: str, kwargs: dict[str, Any]
-) -> SentenceTransformer:
+) -> "SentenceTransformer":
     key = (model, json.dumps(kwargs, sort_keys=True, default=repr))
     with _LOCAL_EMBEDDING_MODELS_LOCK:
         cached = _LOCAL_EMBEDDING_MODELS.get(key)
         if cached is None:
+            from sentence_transformers import SentenceTransformer
+
             cached = SentenceTransformer(model, **kwargs)
             _LOCAL_EMBEDDING_MODELS.clear()
             _LOCAL_EMBEDDING_MODELS[key] = cached

preload.py

@@ -2,62 +2,29 @@ import asyncio
 from helpers import runtime
 from helpers.print_style import PrintStyle
 import models
-from plugins._kokoro_tts.helpers import runtime as kokoro_tts_runtime
-from plugins._whisper_stt.helpers import runtime as whisper_stt_runtime
 
 
 async def preload():
     try:
-        # preload whisper model
-        async def preload_whisper():
-            if not whisper_stt_runtime.is_globally_enabled():
-                return None
-            try:
-                config = whisper_stt_runtime.get_config()
-                return await whisper_stt_runtime.preload(str(config["model_size"]))
-            except Exception as e:
-                PrintStyle().error(f"Error in preload_whisper: {e}")
+        from plugins._model_config.helpers.model_config import get_embedding_model_config_object
 
-        # preload embedding model
-        async def preload_embedding():
-            try:
-                from plugins._model_config.helpers.model_config import get_embedding_model_config_object
-                emb_cfg = get_embedding_model_config_object()
-                if emb_cfg.provider.lower() == "huggingface":
-                    emb_mod = models.get_embedding_model(
-                        emb_cfg.provider,
-                        emb_cfg.name,
-                        model_config=emb_cfg,
-                        **emb_cfg.build_kwargs(),
-                    )
-                    emb_txt = await emb_mod.aembed_query("test")
-                    return emb_txt
-            except Exception as e:
-                PrintStyle().error(f"Error in preload_embedding: {e}")
-
-        # preload kokoro tts model if enabled
-        async def preload_kokoro():
-            if not kokoro_tts_runtime.is_globally_enabled():
-                return None
-            try:
-                return await kokoro_tts_runtime.preload()
-            except Exception as e:
-                PrintStyle().error(f"Error in preload_kokoro: {e}")
-
-        # async tasks to preload
-        tasks = [
-            preload_embedding(),
-            # preload_whisper(),
-            # preload_kokoro()
-        ]
-
-        await asyncio.gather(*tasks, return_exceptions=True)
+        emb_cfg = get_embedding_model_config_object()
+        if emb_cfg.provider.lower() == "huggingface":
+            # Model imports and construction must not block the shared background loop.
+            emb_mod = await asyncio.to_thread(
+                models.get_embedding_model,
+                emb_cfg.provider,
+                emb_cfg.name,
+                model_config=emb_cfg,
+                **emb_cfg.build_kwargs(),
+            )
+            await emb_mod.aembed_query("test")
         PrintStyle().print("Preload completed.")
     except Exception as e:
         PrintStyle().error(f"Error in preload: {e}")
 
 
-# preload transcription model
+# preload the configured embedding model
 if __name__ == "__main__":
     PrintStyle().print("Running preload...")
     runtime.initialize()

run_ui.py

@@ -73,12 +73,12 @@ def create_flush_callback():
 
 @extension.extensible
 def init_a0():
+    initialize.initialize_preload()
     init_chats = initialize.initialize_chats()
     init_chats.result_sync()
 
     initialize.initialize_mcp()
     initialize.initialize_job_loop()
-    initialize.initialize_preload()
 
 
 if __name__ == "__main__":

tests/test_docker_initialize_limits.py

@@ -10,10 +10,10 @@ REPO_ROOT = Path(__file__).resolve().parents[1]
 INITIALIZE_SCRIPT = REPO_ROOT / "docker" / "run" / "fs" / "exe" / "initialize.sh"
 
 
-def _raise_limit_function() -> str:
+def _script_function(name: str) -> str:
     text = INITIALIZE_SCRIPT.read_text(encoding="utf-8")
-    match = re.search(r"^raise_open_file_limit\(\) \{\n.*?^\}\n", text, re.M | re.S)
-    assert match, "initialize.sh must define raise_open_file_limit"
+    match = re.search(rf"^{name}\(\) \{{\n.*?^\}}\n", text, re.M | re.S)
+    assert match, f"initialize.sh must define {name}"
     return match.group(0)
 
 
@@ -27,7 +27,7 @@ def _run_bash(script: str) -> subprocess.CompletedProcess[str]:
 
 
 def test_initialize_raises_soft_open_file_limit_to_requested_target():
-    function = _raise_limit_function()
+    function = _script_function("raise_open_file_limit")
 
     result = _run_bash(
         f"""
@@ -48,7 +48,7 @@ def test_initialize_caps_open_file_limit_at_hard_limit():
     if hard != resource.RLIM_INFINITY and hard < 2048:
         pytest.skip("host hard open-file limit is too low for this regression test")
 
-    function = _raise_limit_function()
+    function = _script_function("raise_open_file_limit")
 
     result = _run_bash(
         f"""
@@ -63,3 +63,29 @@ def test_initialize_caps_open_file_limit_at_hard_limit():
     )
 
     assert "Raised open file soft limit from 1024 to 2048" in result.stdout
+
+
+@pytest.mark.parametrize("package_index", [None, "mirror_Packages", "mirror_Packages.lz4"])
+def test_initialize_updates_only_an_empty_package_cache(tmp_path, package_index):
+    lists = tmp_path / "lists"
+    lists.mkdir()
+    (lists / "partial").mkdir()
+    (lists / "lock").touch()
+    if package_index:
+        (lists / package_index).write_text("cached package index")
+    calls = tmp_path / "apt-calls"
+    function = _script_function("initialize_package_lists").replace(
+        "/var/lib/apt/lists", str(lists)
+    )
+
+    _run_bash(f"""
+        set -euo pipefail
+        apt-get() {{ echo "$*" >> "{calls}"; }}
+        {function}
+        initialize_package_lists
+        wait
+    """)
+
+    assert (calls.read_text() if calls.exists() else "") == (
+        "" if package_index else "update\n"
+    )

tests/test_embedding_model_preload.py

@@ -1,4 +1,8 @@
 from concurrent.futures import ThreadPoolExecutor
+import os
+from pathlib import Path
+import subprocess
+import sys
 import threading
 from types import SimpleNamespace
 
@@ -7,11 +11,60 @@ import pytest
 import models
 
 
+def test_framework_imports_do_not_load_embedding_or_speech_runtimes():
+    result = subprocess.run(
+        [sys.executable, "-c", """
+import sys
+import models
+import preload
+for name in (
+    'sentence_transformers', 'torch', 'whisper',
+    'plugins._kokoro_tts.helpers.runtime',
+    'plugins._whisper_stt.helpers.runtime',
+):
+    assert name not in sys.modules, name
+"""],
+        cwd=Path(__file__).resolve().parents[1],
+        env={**os.environ, "LITELLM_LOCAL_MODEL_COST_MAP": "True"},
+        capture_output=True,
+        text=True,
+        timeout=60,
+    )
+    assert result.returncode == 0, result.stderr
+
+
 def _clear_local_embedding_models():
     with models._LOCAL_EMBEDDING_MODELS_LOCK:
         models._LOCAL_EMBEDDING_MODELS.clear()
 
 
+@pytest.mark.parametrize("skip_init", [False, True])
+def test_runtime_preload_overlaps_chat_restoration_within_extension_hooks(monkeypatch, skip_init):
+    import run_ui
+
+    calls = []
+
+    def extension_hook(point, agent, data):
+        calls.append(point.rsplit("/", 1)[-1])
+        if skip_init and point.endswith("/start"):
+            data["result"] = None
+
+    monkeypatch.setattr(run_ui.extension, "call_extensions_sync", extension_hook)
+    monkeypatch.setattr(run_ui.initialize, "initialize_preload", lambda: calls.append("preload"))
+    monkeypatch.setattr(
+        run_ui.initialize, "initialize_chats",
+        lambda: SimpleNamespace(result_sync=lambda: calls.append("chats")),
+    )
+    monkeypatch.setattr(run_ui.initialize, "initialize_mcp", lambda: calls.append("mcp"))
+    monkeypatch.setattr(run_ui.initialize, "initialize_job_loop", lambda: calls.append("jobs"))
+
+    run_ui.init_a0()
+
+    assert calls == (["start", "end"] if skip_init else [
+        "start", "preload", "chats", "mcp", "jobs", "end",
+    ])
+
+
 def test_local_embedding_preload_is_reused_with_runtime_model_config(monkeypatch):
     created = []
 
@@ -19,7 +72,10 @@ def test_local_embedding_preload_is_reused_with_runtime_model_config(monkeypatch
         def __init__(self, model, **kwargs):
             created.append((model, kwargs))
 
-    monkeypatch.setattr(models, "SentenceTransformer", FakeSentenceTransformer)
+    monkeypatch.setitem(
+        sys.modules, "sentence_transformers",
+        SimpleNamespace(SentenceTransformer=FakeSentenceTransformer),
+    )
     _clear_local_embedding_models()
 
     try:
@@ -63,7 +119,10 @@ def test_local_embedding_cache_tracks_effective_constructor_options(monkeypatch)
         def __init__(self, model, **kwargs):
             created.append((model, kwargs))
 
-    monkeypatch.setattr(models, "SentenceTransformer", FakeSentenceTransformer)
+    monkeypatch.setitem(
+        sys.modules, "sentence_transformers",
+        SimpleNamespace(SentenceTransformer=FakeSentenceTransformer),
+    )
     _clear_local_embedding_models()
 
     try:
@@ -95,7 +154,10 @@ def test_concurrent_preload_and_runtime_share_one_model(monkeypatch):
             construction_started.set()
             assert release_construction.wait(timeout=2)
 
-    monkeypatch.setattr(models, "SentenceTransformer", FakeSentenceTransformer)
+    monkeypatch.setitem(
+        sys.modules, "sentence_transformers",
+        SimpleNamespace(SentenceTransformer=FakeSentenceTransformer),
+    )
     _clear_local_embedding_models()
 
     try:
@@ -130,7 +192,10 @@ def test_failed_model_change_keeps_the_working_cached_model(monkeypatch):
             if model == "broken":
                 raise RuntimeError("model unavailable")
 
-    monkeypatch.setattr(models, "SentenceTransformer", FakeSentenceTransformer)
+    monkeypatch.setitem(
+        sys.modules, "sentence_transformers",
+        SimpleNamespace(SentenceTransformer=FakeSentenceTransformer),
+    )
     _clear_local_embedding_models()
 
     try:
@@ -163,12 +228,14 @@ async def test_preload_uses_the_runtime_embedding_configuration(monkeypatch):
     )
     calls = []
     embedded = []
+    loop_thread = threading.get_ident()
 
     class FakeEmbeddings:
         async def aembed_query(self, text):
             embedded.append(text)
 
     def get_embedding_model(provider, name, **kwargs):
+        assert threading.get_ident() != loop_thread
         calls.append((provider, name, kwargs))
         return FakeEmbeddings()
 
@@ -187,3 +254,34 @@ async def test_preload_uses_the_runtime_embedding_configuration(monkeypatch):
         )
     ]
     assert embedded == ["test"]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("provider", ["openai", "huggingface"])
+async def test_preload_skips_remote_models_and_contains_local_failures(monkeypatch, provider):
+    import preload
+    from plugins._model_config.helpers import model_config
+
+    config = SimpleNamespace(provider=provider, name="example", build_kwargs=lambda: {})
+    calls = []
+    errors = []
+
+    def fail_embedding(*args, **kwargs):
+        calls.append(args)
+        raise RuntimeError("model unavailable")
+
+    monkeypatch.setattr(model_config, "get_embedding_model_config_object", lambda: config)
+    monkeypatch.setattr(models, "get_embedding_model", fail_embedding)
+    monkeypatch.setattr(
+        preload, "PrintStyle",
+        lambda: SimpleNamespace(print=lambda _: None, error=errors.append),
+    )
+
+    await preload.preload()
+
+    if provider == "huggingface":
+        assert calls == [("huggingface", "example")]
+        assert errors == ["Error in preload: model unavailable"]
+    else:
+        assert calls == []
+        assert errors == []