mirror of
https://github.com/agent0ai/agent-zero.git
synced 2026-10-19 15:33:54 +00:00
Speed up chat restoration and bound embedding warm-up
Reuse Git version metadata within each chat-load operation and pass saved message token counts into reconstruction. Preserve per-agent settings and credential reads, initialization hooks, profiles and legacy token behavior. Give the existing deferred embedding preload up to 10 extra seconds after chat restoration, within extensible init_a0. Continue startup on timeout or task failure while allowing unfinished model loading to complete. Across five trials per version with the same chat workload and a cached embedding model, median HTTP readiness after the restart command returned fell from 38.1 to 27.8 seconds: 10.3 seconds saved, or 27% less waiting. Boot-to-first-token fell from 42.7 to 33.4 seconds. Post-send first-token latency did not improve in this sample (4.709 to 5.108 seconds). Validation: 146 focused framework tests passed, followed by 20 preload tests including seven new edge cases. Real ten-second stalled-loading/inference checks and offline missing/cached-model checks passed. Five DeferredTask lifecycle tests passed separately; an existing test-order failure when combined with preload imports also reproduces without the new cases. Fresh-install performance was not measured; the full repository suite was not run.
This commit is contained in:
parent
2afffe8edc
commit
cc098d3e98
12 changed files with 363 additions and 29 deletions
|
|
@ -22,7 +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.
|
||||
- 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 and then allows up to ten seconds for it to finish before later startup hooks compete for CPU. This caps only the extra preload wait; chat restoration and extension hooks are outside that deadline. Timeout leaves preload running, so a first chat can still wait for an uncached model. Waiting is best-effort: timeout or task failure does not prevent HTTP startup, and extensions may disable preload. Initialization extension hooks remain in control. `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.
|
||||
|
|
|
|||
|
|
@ -161,6 +161,7 @@ class Message(Record):
|
|||
msg = Message(
|
||||
ai=data["ai"],
|
||||
content=content,
|
||||
tokens=data.get("tokens", 0),
|
||||
id=data.get("id", ""),
|
||||
metadata=metadata,
|
||||
sequence=int(data.get("sequence", 0) or 0),
|
||||
|
|
|
|||
|
|
@ -80,6 +80,7 @@
|
|||
- Update this file whenever public functions, classes, persistence behavior, path/security assumptions, side effects, or cross-module contracts change.
|
||||
- `clear_responses_provider_state(agent)` removes the active provider continuation IDs after local history rewrites while preserving stored response ID lists for later cleanup.
|
||||
- `Message.from_dict()` normalizes legacy AI Responses metadata through `LLMResult.metadata()` so loaded chats shed transient payloads while unrelated metadata and non-AI tool-result inputs remain intact.
|
||||
- `Message.from_dict()` passes saved token counts to the constructor to avoid recounting restored messages. It still restores the saved value after the summary, preserving zero/missing-count lazy recalculation against the summarized content.
|
||||
- Topic token totals and large-message compression include retained extras snapshots when the selected model uses Responses. Chat Completions keeps text-only accounting. Summarized messages/topics no longer contribute snapshot overhead; unchanged records conservatively count it even when replay eligibility later changes.
|
||||
- `output_langchain()` removes leading assistant messages after grouping so provider histories always begin with a user turn; the WebUI greeting remains persisted and displayed but is not sent as an orphaned assistant message.
|
||||
- `_json_dumps()` emits compact JSON (`","`, `":"` separators) for serialized history and generated non-string user-turn content.
|
||||
|
|
|
|||
|
|
@ -7,7 +7,7 @@ from datetime import datetime
|
|||
from typing import Any
|
||||
|
||||
from agent import Agent, AgentConfig, AgentContext, AgentContextType
|
||||
from helpers import files, history
|
||||
from helpers import files, history, settings
|
||||
from helpers.litellm_transport import delete_stored_response_ids
|
||||
from helpers.localization import Localization
|
||||
from initialize import initialize_agent
|
||||
|
|
@ -80,15 +80,16 @@ def load_tmp_chats():
|
|||
json_files.append(chat_file)
|
||||
|
||||
ctxids = []
|
||||
for file in json_files:
|
||||
try:
|
||||
js = files.read_file(file)
|
||||
data = json.loads(js)
|
||||
ctx = _deserialize_context(data)
|
||||
mark_chat_saved(ctx)
|
||||
ctxids.append(ctx.id)
|
||||
except Exception as e:
|
||||
print(f"Error loading chat {file}: {e}")
|
||||
with settings.version_cache_scope():
|
||||
for file in json_files:
|
||||
try:
|
||||
js = files.read_file(file)
|
||||
data = json.loads(js)
|
||||
ctx = _deserialize_context(data)
|
||||
mark_chat_saved(ctx)
|
||||
ctxids.append(ctx.id)
|
||||
except Exception as e:
|
||||
print(f"Error loading chat {file}: {e}")
|
||||
return ctxids
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -45,6 +45,7 @@
|
|||
- Serialized chats store `agent_profile` both at the context level for the main chat and on each serialized agent so subordinate profiles survive server restart.
|
||||
- Deserialization must rebuild each agent with its serialized profile when present, falling back to the context profile for older chat files.
|
||||
- Chat loading skips directories that do not contain `chat.json`; malformed existing chat files still report load errors.
|
||||
- `load_tmp_chats()` reuses one Git version through `settings.version_cache_scope()` for its per-file restoration loop. Each agent still receives normal settings/credential reads and initialization hooks; the scope expires after the load, including failures.
|
||||
- Chat saves write and fsync a same-directory temporary file, atomically replace `chat.json`, and fsync the directory so an interrupted save cannot truncate the previous chat.
|
||||
- Contexts are marked with private `SAVED_CHAT_CONTEXT_DATA_KEY` only after a successful save or load from disk so snapshot code can detect deleted chat files without hiding fresh unsaved chats.
|
||||
- Chat and message-file removal reject empty or whitespace-only context IDs before provider or filesystem deletion begins.
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
import base64
|
||||
from contextlib import contextmanager
|
||||
from contextvars import ContextVar, Token
|
||||
from copy import deepcopy
|
||||
import hashlib
|
||||
|
|
@ -181,8 +182,8 @@ UI_CONTROL_VISIBILITY_DEFAULTS = {
|
|||
|
||||
SETTINGS_FILE = files.get_abs_path("usr/settings.json")
|
||||
_settings: Settings | None = None
|
||||
_apply_settings_cache: ContextVar[dict[str, str] | None] = ContextVar(
|
||||
"apply_settings_cache", default=None
|
||||
_settings_version_cache: ContextVar[dict[str, str] | None] = ContextVar(
|
||||
"settings_version_cache", default=None
|
||||
)
|
||||
_runtime_settings_snapshot: Settings | None = None
|
||||
_prompt_settings_snapshot: ContextVar[Settings | None] = ContextVar(
|
||||
|
|
@ -421,20 +422,26 @@ def set_runtime_settings_snapshot(settings: Settings) -> None:
|
|||
_runtime_settings_snapshot = settings.copy()
|
||||
|
||||
|
||||
@contextmanager
|
||||
def version_cache_scope(version: str | None = None):
|
||||
cached = {"version": version or _get_version()}
|
||||
token = _settings_version_cache.set(cached)
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
# Expire the cache in inherited task contexts too.
|
||||
cached.clear()
|
||||
_settings_version_cache.reset(token)
|
||||
|
||||
|
||||
def set_settings(settings: Settings, apply: bool = True, browser_timezone: str | None = None):
|
||||
global _settings
|
||||
previous = _settings
|
||||
_settings = normalize_settings(settings)
|
||||
_write_settings_file(_settings)
|
||||
if apply:
|
||||
cached = {"version": _settings["version"]}
|
||||
token = _apply_settings_cache.set(cached)
|
||||
try:
|
||||
with version_cache_scope(_settings["version"]):
|
||||
_apply_settings(previous, browser_timezone)
|
||||
finally:
|
||||
# Expire the cache in inherited task contexts too.
|
||||
cached.clear()
|
||||
_apply_settings_cache.reset(token)
|
||||
return reload_settings()
|
||||
|
||||
|
||||
|
|
@ -853,4 +860,4 @@ def create_auth_token() -> str:
|
|||
|
||||
|
||||
def _get_version():
|
||||
return (_apply_settings_cache.get() or {}).get("version") or git.get_version()
|
||||
return (_settings_version_cache.get() or {}).get("version") or git.get_version()
|
||||
|
|
|
|||
|
|
@ -34,6 +34,7 @@
|
|||
- `get_settings() -> Settings`
|
||||
- `reload_settings() -> Settings`
|
||||
- `set_runtime_settings_snapshot(settings: Settings) -> None`
|
||||
- `version_cache_scope(version: str | None = None)`: Reuse only the Git version within one synchronous operation, expiring inherited copies on exit.
|
||||
- `set_settings(settings: Settings, apply: bool=..., browser_timezone: str | None=...)`
|
||||
- `set_settings_delta(delta: dict, apply: bool=...)`
|
||||
- `merge_settings(original: Settings, delta: dict) -> Settings`
|
||||
|
|
@ -70,11 +71,12 @@
|
|||
|
||||
- Important called helpers/classes observed in the source: `TypeVar`, `files.get_abs_path`, `dotenv.get_dotenv_value`, `opts.insert`, `str.strip`, `_is_valid_timezone`, `str.strip.lower`, `_normalize_timezone_setting`, `SettingsOutput`, `get_default_settings`, `_ensure_option_present`, `_resolve_runtime_timezone`, `get_default_secrets_manager`, `get_settings`, `get_settings_for_prompt`, `normalize_settings`, `_load_sensitive_settings`, `deepcopy`, `settings.copy`, `_write_settings_file`, `reload_settings`, `set_settings`, `initialize_agent`.
|
||||
- Applying settings refreshes active context configs while preserving each subordinate agent's own profile.
|
||||
- Agent refreshes reuse the version already resolved by `set_settings()` in a
|
||||
task-local scope, restored even on failure. Settings normalization, credential
|
||||
reads, and initialization hooks still run for each agent; the final reload and
|
||||
subsequent saves resolve the version afresh. The cache is cleared on exit so
|
||||
deferred tasks cannot retain an expired version through an inherited context.
|
||||
- `version_cache_scope()` eagerly resolves or accepts a Git version in a task-local
|
||||
scope. Settings application and persisted-chat loading reuse it while retaining
|
||||
per-agent normalization, credential reads, and initialization hooks. Nested scopes
|
||||
restore their parent even on failure. Each scope clears its cache on exit so
|
||||
inherited task contexts cannot retain or repopulate an expired version; reads
|
||||
outside an active scope resolve the version afresh.
|
||||
- Applying settings starts a deferred `MCPConfig.update(...)` with the current `mcp_servers` string when global MCP server settings change.
|
||||
- `get_settings()` retains normalize-on-read behavior. Prompt-building callers
|
||||
explicitly use `get_settings_for_prompt()` to reuse one task-local snapshot
|
||||
|
|
|
|||
11
run_ui.py
11
run_ui.py
|
|
@ -73,10 +73,19 @@ def create_flush_callback():
|
|||
|
||||
@extension.extensible
|
||||
def init_a0():
|
||||
initialize.initialize_preload()
|
||||
init_preload = initialize.initialize_preload()
|
||||
init_chats = initialize.initialize_chats()
|
||||
init_chats.result_sync()
|
||||
|
||||
# Give the first chat a warm model without waiting indefinitely on downloads.
|
||||
if init_preload:
|
||||
try:
|
||||
init_preload.result_sync(timeout=10)
|
||||
except TimeoutError:
|
||||
pass
|
||||
except Exception as e:
|
||||
PrintStyle().error(f"Error waiting for preload: {e}")
|
||||
|
||||
initialize.initialize_mcp()
|
||||
initialize.initialize_job_loop()
|
||||
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ from pathlib import Path
|
|||
import subprocess
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
|
|
@ -50,7 +51,16 @@ def test_runtime_preload_overlaps_chat_restoration_within_extension_hooks(monkey
|
|||
data["result"] = None
|
||||
|
||||
monkeypatch.setattr(run_ui.extension, "call_extensions_sync", extension_hook)
|
||||
monkeypatch.setattr(run_ui.initialize, "initialize_preload", lambda: calls.append("preload"))
|
||||
def preload():
|
||||
calls.append("preload")
|
||||
|
||||
def wait(timeout):
|
||||
assert timeout == 10
|
||||
calls.append("preload_ready")
|
||||
|
||||
return SimpleNamespace(result_sync=wait)
|
||||
|
||||
monkeypatch.setattr(run_ui.initialize, "initialize_preload", preload)
|
||||
monkeypatch.setattr(
|
||||
run_ui.initialize, "initialize_chats",
|
||||
lambda: SimpleNamespace(result_sync=lambda: calls.append("chats")),
|
||||
|
|
@ -61,10 +71,199 @@ def test_runtime_preload_overlaps_chat_restoration_within_extension_hooks(monkey
|
|||
run_ui.init_a0()
|
||||
|
||||
assert calls == (["start", "end"] if skip_init else [
|
||||
"start", "preload", "chats", "mcp", "jobs", "end",
|
||||
"start", "preload", "chats", "preload_ready", "mcp", "jobs", "end",
|
||||
])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("error", [None, TimeoutError(), RuntimeError("preload failed")])
|
||||
def test_runtime_preload_wait_remains_best_effort(monkeypatch, error):
|
||||
import run_ui
|
||||
|
||||
calls = []
|
||||
errors = []
|
||||
|
||||
def wait(timeout):
|
||||
assert timeout == 10
|
||||
raise error
|
||||
|
||||
monkeypatch.setattr(run_ui.extension, "call_extensions_sync", lambda *_args, **_kwargs: None)
|
||||
monkeypatch.setattr(
|
||||
run_ui.initialize, "initialize_preload",
|
||||
lambda: SimpleNamespace(result_sync=wait) if error is not None else None,
|
||||
)
|
||||
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"))
|
||||
monkeypatch.setattr(run_ui, "PrintStyle", lambda: SimpleNamespace(error=errors.append))
|
||||
|
||||
run_ui.init_a0()
|
||||
|
||||
assert calls == ["chats", "mcp", "jobs"]
|
||||
assert errors == (["Error waiting for preload: preload failed"] if isinstance(error, RuntimeError) else [])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("blocked_stage", ["construction", "encoding"])
|
||||
def test_startup_continues_after_ten_seconds_with_model_still_loading(monkeypatch, blocked_stage, record_property):
|
||||
import preload
|
||||
import run_ui
|
||||
from helpers import persist_chat
|
||||
from helpers.defer import EventLoopThread
|
||||
from plugins._model_config.helpers import model_config
|
||||
|
||||
release = threading.Event()
|
||||
blocked = threading.Event()
|
||||
chats_loaded = threading.Event()
|
||||
warmed = threading.Event()
|
||||
preload_done = threading.Event()
|
||||
later_startup = []
|
||||
errors = []
|
||||
constructions = []
|
||||
config = models.ModelConfig(models.ModelType.EMBEDDING, "huggingface", "sentence-transformers/timeout-test")
|
||||
|
||||
def extension_hook(point, agent, data):
|
||||
if point.endswith("/initialize_preload/end"):
|
||||
data["result"].add_done_callback(lambda _: preload_done.set())
|
||||
|
||||
def hold():
|
||||
blocked.set()
|
||||
assert release.wait(20), "Test did not release the model operation"
|
||||
|
||||
class HeldModel:
|
||||
def __init__(self, name, **kwargs):
|
||||
constructions.append(name)
|
||||
if blocked_stage == "construction":
|
||||
hold()
|
||||
|
||||
def encode(self, inputs, **kwargs):
|
||||
if blocked_stage == "encoding":
|
||||
hold()
|
||||
warmed.set()
|
||||
return [[1.0] for _ in inputs]
|
||||
|
||||
monkeypatch.setitem(sys.modules, "sentence_transformers", SimpleNamespace(SentenceTransformer=HeldModel))
|
||||
monkeypatch.setattr(model_config, "get_embedding_model_config_object", lambda: config)
|
||||
monkeypatch.setattr(models, "get_provider_config", lambda *_args: None)
|
||||
monkeypatch.setattr(models, "get_api_key_raw", lambda *_args: "None")
|
||||
monkeypatch.setattr(run_ui.extension, "call_extensions_sync", extension_hook)
|
||||
monkeypatch.setattr(persist_chat, "load_tmp_chats", chats_loaded.set)
|
||||
monkeypatch.setattr(run_ui.initialize, "initialize_mcp", lambda: later_startup.append("mcp"))
|
||||
monkeypatch.setattr(run_ui.initialize, "initialize_job_loop", lambda: later_startup.append("jobs"))
|
||||
monkeypatch.setattr(preload, "PrintStyle", lambda: SimpleNamespace(print=lambda _: None, error=errors.append))
|
||||
monkeypatch.setattr(run_ui, "PrintStyle", lambda: SimpleNamespace(error=errors.append))
|
||||
_clear_local_embedding_models()
|
||||
|
||||
# Release blocked work even if a regression removes the startup timeout.
|
||||
with ThreadPoolExecutor(max_workers=1) as pool:
|
||||
started = time.monotonic()
|
||||
startup = pool.submit(run_ui.init_a0)
|
||||
try:
|
||||
assert blocked.wait(3)
|
||||
assert chats_loaded.wait(3)
|
||||
startup.result(timeout=12)
|
||||
elapsed = time.monotonic() - started
|
||||
record_property("startup_wait_seconds", round(elapsed, 3))
|
||||
assert 9.5 <= elapsed < 13, elapsed
|
||||
assert later_startup == ["mcp", "jobs"]
|
||||
assert not warmed.is_set()
|
||||
assert not release.is_set()
|
||||
|
||||
release.set()
|
||||
assert preload_done.wait(3)
|
||||
assert warmed.is_set()
|
||||
runtime_model = models.get_embedding_model(config.provider, config.name, model_config=config)
|
||||
assert runtime_model.embed_query("first message") == [1.0]
|
||||
assert constructions == ["timeout-test"]
|
||||
assert not errors
|
||||
finally:
|
||||
release.set()
|
||||
startup.result(timeout=12)
|
||||
assert preload_done.wait(3)
|
||||
EventLoopThread().terminate()
|
||||
_clear_local_embedding_models()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("scenario", ["cached", "missing", "broken", "remote", "busy_restoration"])
|
||||
def test_real_preload_startup_handles_ready_or_unavailable_models(monkeypatch, scenario, record_property):
|
||||
import preload
|
||||
import run_ui
|
||||
from helpers import persist_chat
|
||||
from helpers.defer import EventLoopThread
|
||||
from plugins._model_config.helpers import model_config
|
||||
|
||||
chats_started = threading.Event()
|
||||
release_chats = threading.Event()
|
||||
model_constructed = threading.Event()
|
||||
preload_done = threading.Event()
|
||||
constructions = []
|
||||
errors = []
|
||||
calls = []
|
||||
config = models.ModelConfig(
|
||||
models.ModelType.EMBEDDING,
|
||||
"openai" if scenario == "remote" else "huggingface",
|
||||
"sentence-transformers/startup-test",
|
||||
)
|
||||
|
||||
class Model:
|
||||
def __init__(self, name, **kwargs):
|
||||
constructions.append(name)
|
||||
if scenario in {"missing", "broken"}:
|
||||
raise OSError(scenario + " model")
|
||||
model_constructed.set()
|
||||
|
||||
def encode(self, inputs, **kwargs):
|
||||
return [[1.0] for _ in inputs]
|
||||
|
||||
def restore_chats():
|
||||
chats_started.set()
|
||||
if scenario == "busy_restoration":
|
||||
assert release_chats.wait(5)
|
||||
|
||||
def extension_hook(point, agent, data):
|
||||
if point.endswith("/initialize_preload/end"):
|
||||
data["result"].add_done_callback(lambda _: preload_done.set())
|
||||
|
||||
monkeypatch.setitem(sys.modules, "sentence_transformers", SimpleNamespace(SentenceTransformer=Model))
|
||||
monkeypatch.setattr(models, "get_provider_config", lambda *_args: None)
|
||||
monkeypatch.setattr(models, "get_api_key_raw", lambda *_args: "None")
|
||||
monkeypatch.setattr(model_config, "get_embedding_model_config_object", lambda: config)
|
||||
monkeypatch.setattr(run_ui.extension, "call_extensions_sync", extension_hook)
|
||||
monkeypatch.setattr(persist_chat, "load_tmp_chats", restore_chats)
|
||||
monkeypatch.setattr(run_ui.initialize, "initialize_mcp", lambda: calls.append("mcp"))
|
||||
monkeypatch.setattr(run_ui.initialize, "initialize_job_loop", lambda: calls.append("jobs"))
|
||||
monkeypatch.setattr(preload, "PrintStyle", lambda: SimpleNamespace(print=lambda _: None, error=errors.append))
|
||||
monkeypatch.setattr(run_ui, "PrintStyle", lambda: SimpleNamespace(error=errors.append))
|
||||
_clear_local_embedding_models()
|
||||
if scenario == "cached":
|
||||
models.get_embedding_model(config.provider, config.name)
|
||||
|
||||
with ThreadPoolExecutor(max_workers=1) as pool:
|
||||
started = time.monotonic()
|
||||
startup = pool.submit(run_ui.init_a0)
|
||||
try:
|
||||
assert chats_started.wait(2)
|
||||
if scenario == "busy_restoration":
|
||||
assert model_constructed.wait(2)
|
||||
assert not startup.done()
|
||||
release_chats.set()
|
||||
startup.result(timeout=2)
|
||||
elapsed = time.monotonic() - started
|
||||
record_property("startup_wait_seconds", round(elapsed, 3))
|
||||
assert elapsed < 2, elapsed
|
||||
assert preload_done.wait(2)
|
||||
assert calls == ["mcp", "jobs"]
|
||||
assert len(constructions) == (0 if scenario == "remote" else 1)
|
||||
assert errors == (["Error in preload: " + scenario + " model"] if scenario in {"missing", "broken"} else [])
|
||||
finally:
|
||||
release_chats.set()
|
||||
startup.result(timeout=12)
|
||||
assert preload_done.wait(2)
|
||||
EventLoopThread().terminate()
|
||||
_clear_local_embedding_models()
|
||||
|
||||
|
||||
def test_local_embedding_preload_is_reused_with_runtime_model_config(monkeypatch):
|
||||
created = []
|
||||
|
||||
|
|
|
|||
|
|
@ -1,3 +1,6 @@
|
|||
import pytest
|
||||
|
||||
from helpers import history
|
||||
from helpers.history import output_langchain
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
|
||||
|
|
@ -12,3 +15,42 @@ def test_output_langchain_omits_leading_assistant_messages():
|
|||
)
|
||||
|
||||
assert messages == [HumanMessage("Hello"), AIMessage("Hi")]
|
||||
|
||||
|
||||
def test_restored_message_reuses_saved_tokens_and_summary(monkeypatch):
|
||||
data = {
|
||||
"_cls": "Message", "ai": False, "content": "original content",
|
||||
"id": "message", "sequence": 2, "metadata": {"custom": True},
|
||||
"summary": "saved summary", "tokens": 42,
|
||||
}
|
||||
|
||||
def recount(_text):
|
||||
pytest.fail("Restoring a saved count must not tokenize the message again")
|
||||
|
||||
monkeypatch.setattr(history.tokens, "approximate_tokens", recount)
|
||||
restored = history.Message.from_dict(data, history=None)
|
||||
|
||||
assert restored.to_dict() == data
|
||||
assert restored.get_tokens() == 42
|
||||
assert restored.output()[0]["content"] == "saved summary"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("saved", [{}, {"tokens": 0}, {"tokens": None}])
|
||||
def test_restored_message_without_token_count_recounts_using_summary(monkeypatch, saved):
|
||||
counted = []
|
||||
|
||||
def count(text):
|
||||
counted.append(text)
|
||||
return 7
|
||||
|
||||
monkeypatch.setattr(history.tokens, "approximate_tokens", count)
|
||||
restored = history.Message.from_dict(
|
||||
{"ai": False, "content": "original", "summary": "summary", **saved},
|
||||
history=None,
|
||||
)
|
||||
|
||||
assert restored.tokens == saved.get("tokens", 0)
|
||||
assert restored.get_tokens() == 7
|
||||
assert counted == ["user: original", "user: summary"]
|
||||
assert restored.get_tokens() == 7
|
||||
assert len(counted) == 2
|
||||
|
|
|
|||
|
|
@ -58,6 +58,48 @@ def test_load_tmp_chats_skips_directories_without_chat_json(monkeypatch, capsys)
|
|||
assert "Error loading chat" not in capsys.readouterr().out
|
||||
|
||||
|
||||
def test_load_tmp_chats_reuses_only_version_and_isolates_bad_files(
|
||||
monkeypatch, tmp_path, capsys
|
||||
) -> None:
|
||||
from helpers import persist_chat, settings
|
||||
|
||||
configured = settings.get_default_settings()
|
||||
version_reads = []
|
||||
sensitive_reads = []
|
||||
versions = []
|
||||
|
||||
def get_version():
|
||||
version_reads.append(len(version_reads))
|
||||
return str(len(version_reads))
|
||||
|
||||
def deserialize(data):
|
||||
versions.append(settings.get_settings()["version"])
|
||||
return SimpleNamespace(id=data["id"], data={})
|
||||
|
||||
monkeypatch.setattr(settings, "_settings", configured)
|
||||
monkeypatch.setattr(settings.git, "get_version", get_version)
|
||||
monkeypatch.setattr(settings, "_load_sensitive_settings", sensitive_reads.append)
|
||||
monkeypatch.setattr(persist_chat, "CHATS_FOLDER", str(tmp_path))
|
||||
monkeypatch.setattr(persist_chat, "_deserialize_context", deserialize)
|
||||
monkeypatch.setattr(persist_chat.files, "list_files", lambda *_: ["first", "bad", "last"])
|
||||
monkeypatch.setattr(persist_chat, "_convert_v080_chats", lambda: None)
|
||||
for name in ("first", "bad", "last"):
|
||||
folder = tmp_path / name
|
||||
folder.mkdir()
|
||||
(folder / "chat.json").write_text("{broken" if name == "bad" else json.dumps({"id": name}))
|
||||
|
||||
assert persist_chat.load_tmp_chats() == ["first", "last"]
|
||||
assert versions == ["1", "1"]
|
||||
assert len(sensitive_reads) == 2
|
||||
assert "Error loading chat" in capsys.readouterr().out
|
||||
assert settings._get_version() == "2"
|
||||
|
||||
assert persist_chat.load_tmp_chats() == ["first", "last"]
|
||||
assert versions == ["1", "1", "3", "3"]
|
||||
assert len(sensitive_reads) == 4
|
||||
assert settings._get_version() == "4"
|
||||
|
||||
|
||||
def test_save_tmp_chat_preserves_existing_file_until_atomic_replace(
|
||||
monkeypatch, tmp_path
|
||||
) -> None:
|
||||
|
|
|
|||
|
|
@ -6,6 +6,35 @@ import pytest
|
|||
from helpers import extension, settings
|
||||
|
||||
|
||||
def test_version_cache_scope_nests_and_expires_in_inherited_contexts(monkeypatch):
|
||||
reads = []
|
||||
|
||||
def get_version():
|
||||
reads.append("fresh")
|
||||
return "fresh"
|
||||
|
||||
monkeypatch.setattr(settings.git, "get_version", get_version)
|
||||
|
||||
with settings.version_cache_scope("outer"):
|
||||
outer_context = copy_context()
|
||||
with pytest.raises(RuntimeError, match="scope failed"):
|
||||
with settings.version_cache_scope():
|
||||
assert settings._get_version() == "outer"
|
||||
with settings.version_cache_scope("inner"):
|
||||
assert settings._get_version() == "inner"
|
||||
inner_context = copy_context()
|
||||
assert settings._get_version() == "outer"
|
||||
raise RuntimeError("scope failed")
|
||||
assert settings._get_version() == "outer"
|
||||
assert Context().run(settings._get_version) == "fresh"
|
||||
|
||||
assert settings._get_version() == "fresh"
|
||||
assert outer_context.run(settings._get_version) == "fresh"
|
||||
assert inner_context.run(settings._get_version) == "fresh"
|
||||
assert inner_context.run(settings._get_version) == "fresh"
|
||||
assert reads == ["fresh"] * 5
|
||||
|
||||
|
||||
@pytest.mark.parametrize("fail_apply", [False, True])
|
||||
def test_settings_apply_reuses_only_version_and_restores_scope(monkeypatch, fail_apply):
|
||||
configured = settings.get_default_settings()
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue