diff --git a/Cortex_Preview.py b/Cortex_Preview.py index 78c944e..e397316 100644 --- a/Cortex_Preview.py +++ b/Cortex_Preview.py @@ -11,6 +11,7 @@ ROOT = Path(__file__).resolve().parent sys.path.insert(0, str(ROOT / "backend")) +import httpx # noqa: E402 import ollama # noqa: E402 import uvicorn # noqa: E402 @@ -84,11 +85,17 @@ def build_preview_app( database.migrate_from_json_if_needed() permanent_memory = PermanentMemoryManager(app_paths=paths) settings_repository = SQLiteSettingsRepository( - paths.database, + paths.settings_database, legacy=LegacySettingsReader(), + # Settings used to live inside the chat database; adopt them once so + # an upgrade does not silently revert to defaults. + adopt_from=paths.database, ) ollama_host = os.environ.get("CORTEX_OLLAMA_HOST", "http://127.0.0.1:11434") - client = ollama.Client(host=ollama_host) + client = ollama.Client( + host=ollama_host, + timeout=httpx.Timeout(connect=5.0, read=600.0, write=30.0, pool=5.0), + ) def gguf_directory() -> Path: # Re-read settings each call (cheap SQLite read, same pattern the API diff --git a/backend/cortex_backend/api/app.py b/backend/cortex_backend/api/app.py index 1cf2f53..1734795 100644 --- a/backend/cortex_backend/api/app.py +++ b/backend/cortex_backend/api/app.py @@ -5,6 +5,7 @@ from contextlib import asynccontextmanager from dataclasses import dataclass from datetime import datetime, timezone +import logging import os import tempfile from pathlib import Path @@ -42,6 +43,8 @@ from .routes import build_router from .security import SessionManager +logger = logging.getLogger(__name__) + @dataclass(slots=True) class BackendDependencies: @@ -159,7 +162,16 @@ async def lifespan(app: FastAPI): yield finally: app.state.ready = False - await app.state.jobs.shutdown() + # Runtime teardown below must run even if job shutdown raises -- + # it is what actually terminates the llama-server child process, + # and a worker that refuses to unwind must not leave that + # process (and the GPU/RAM it holds) orphaned. + try: + await app.state.jobs.shutdown() + except Exception: + logger.exception( + "Cortex job registry shutdown raised; continuing with runtime teardown." + ) if app.state.execution_lifecycle is not None: app.state.execution_lifecycle.stop() elif app.state.execution_coordinator is not None: diff --git a/backend/cortex_backend/api/jobs.py b/backend/cortex_backend/api/jobs.py index c8ab408..c5b2729 100644 --- a/backend/cortex_backend/api/jobs.py +++ b/backend/cortex_backend/api/jobs.py @@ -197,13 +197,22 @@ class JobRegistry: registry itself remains the single authority for lifecycle transitions. """ - def __init__(self, *, poll_seconds: float = 0.025, max_terminal_jobs: int = 100): + def __init__( + self, + *, + poll_seconds: float = 0.025, + max_terminal_jobs: int = 100, + shutdown_grace_seconds: float = 10.0, + ): if poll_seconds <= 0: raise ValueError("poll_seconds must be positive") if max_terminal_jobs <= 0: raise ValueError("max_terminal_jobs must be positive") + if shutdown_grace_seconds <= 0: + raise ValueError("shutdown_grace_seconds must be positive") self._poll_seconds = poll_seconds self._max_terminal_jobs = max_terminal_jobs + self._shutdown_grace_seconds = shutdown_grace_seconds self._records: dict[str, _JobRecord] = {} self._active: dict[JobKind, str] = {} self._request_index: dict[tuple[JobKind, str, str], str] = {} @@ -472,7 +481,23 @@ async def events( await asyncio.sleep(self._poll_seconds) async def shutdown(self) -> None: - """Request cancellation and wait for owned workers to finish safely.""" + """Request cancellation and wait for owned workers to finish safely. + + A worker that has already begun committing its result (see + :meth:`JobProgressSink.begin_commit`) is awaited without a bound -- + that commit must finish so persisted state and the retained event + stream stay consistent. A worker that has not committed is only + cooperative on a best-effort basis: it may be blocked inside a + synchronous call (a model HTTP request with no read deadline, for + example) that never polls ``cancel_event``. Waiting on it + indefinitely would hang app shutdown -- and the llama-server child + process it is talking to -- for as long as that call takes, so the + first wait below is capped at ``shutdown_grace_seconds``. Anything + still pending after that grace period is re-checked: a worker that + committed *during* the grace period still gets the unbounded wait + it is owed; a worker that never committed is abandoned so shutdown + can proceed with the rest of teardown. + """ with self._lock: self._accepting = False records = [ @@ -498,12 +523,31 @@ async def shutdown(self) -> None: if self._active.get(record.kind) == record.job_id: self._active.pop(record.kind, None) tasks = [record.task for record in self._records.values() if record.task] - pending = [task for task in tasks if task is not asyncio.current_task()] - if pending: + pending = {task for task in tasks if task is not asyncio.current_task()} + if not pending: + return + _, still_pending = await asyncio.wait(pending, timeout=self._shutdown_grace_seconds) + if not still_pending: + return + with self._lock: + committed_still_pending = [ + record.task + for record in self._records.values() + if record.task in still_pending and record.commit_started + ] + abandoned = len(still_pending) - len(committed_still_pending) + if abandoned: + logging.warning( + "Cortex shutdown: %d job worker(s) did not observe cancellation within " + "%.0fs and were abandoned so shutdown could proceed.", + abandoned, + self._shutdown_grace_seconds, + ) + if committed_still_pending: # Cancelling an asyncio task does not stop its ``to_thread`` # worker. Waiting for the task lets the worker observe the event, # complete its cleanup, and finalize the cancellation itself. - await asyncio.gather(*pending, return_exceptions=True) + await asyncio.gather(*committed_still_pending, return_exceptions=True) async def _run( self, diff --git a/backend/cortex_backend/api/routes.py b/backend/cortex_backend/api/routes.py index 815df47..30630b5 100644 --- a/backend/cortex_backend/api/routes.py +++ b/backend/cortex_backend/api/routes.py @@ -200,7 +200,9 @@ def handoff(request: Request) -> HandoffResponse: manager.validate_request_context(request) supplied = request.headers.get("X-Cortex-Handoff", "") expected = request.app.state.handoff_secret - if not expected or not hmac.compare_digest(supplied, expected): + if not expected or not hmac.compare_digest( + supplied.encode("latin-1"), expected.encode("utf-8") + ): raise HTTPException(status_code=401, detail="Cortex handoff unavailable.") token, expires_at = manager.issue_bootstrap_token() return HandoffResponse(bootstrap_token=token, expires_at=expires_at) diff --git a/backend/cortex_backend/api/security.py b/backend/cortex_backend/api/security.py index b0ede32..eaaf57b 100644 --- a/backend/cortex_backend/api/security.py +++ b/backend/cortex_backend/api/security.py @@ -2,7 +2,7 @@ from __future__ import annotations -from dataclasses import dataclass +from dataclasses import dataclass, replace from datetime import datetime, timedelta, timezone import hashlib import hmac @@ -34,6 +34,10 @@ class SessionPrincipal: session_id: str installation_principal_id: str expires_at: datetime + # When this session was first issued -- fixed for its whole life, unlike + # expires_at, so a sliding renewal has an absolute lifetime to cap + # against instead of extending forever under continuous use. + issued_at: datetime class SessionManager: @@ -44,15 +48,19 @@ def __init__( *, bootstrap_token: str | None = None, ttl_seconds: int = 3600, + max_lifetime_seconds: int = 24 * 3600, allowed_hosts: Iterable[str] = ("127.0.0.1", "localhost", "::1"), installation_principal_id: str | None = None, ): if ttl_seconds < 60: raise ValueError("session TTL must be at least 60 seconds") + if max_lifetime_seconds < ttl_seconds: + raise ValueError("max_lifetime_seconds must be at least ttl_seconds") self._bootstrap_token = bootstrap_token or secrets.token_urlsafe(32) self._bootstrap_expires_at = datetime.now(timezone.utc) + timedelta(minutes=5) self._bootstrap_used = False self._ttl = timedelta(seconds=ttl_seconds) + self._max_lifetime = timedelta(seconds=max_lifetime_seconds) self._allowed_hosts = frozenset(host.lower() for host in allowed_hosts) self._installation_principal_id = validate_installation_principal_id( installation_principal_id or secrets.token_hex(32) @@ -89,30 +97,43 @@ def installation_principal_id(self) -> str: def exchange(self, bootstrap_token: str) -> SessionExchange: with self._lock: if self._bootstrap_used or self._bootstrap_expires_at <= datetime.now(timezone.utc) or not hmac.compare_digest( - bootstrap_token, - self._bootstrap_token, + bootstrap_token.encode("utf-8"), + self._bootstrap_token.encode("utf-8"), ): raise SessionSecurityError("invalid bootstrap token") self._bootstrap_used = True raw_token = secrets.token_urlsafe(32) - expires_at = datetime.now(timezone.utc) + self._ttl + issued_at = datetime.now(timezone.utc) session_id = secrets.token_urlsafe(16) - self._sessions[self._digest(raw_token)] = SessionPrincipal( + principal = SessionPrincipal( session_id=session_id, installation_principal_id=self._installation_principal_id, - expires_at=expires_at, - ) - return SessionExchange( - token=raw_token, - principal=self._sessions[self._digest(raw_token)], + expires_at=issued_at + self._ttl, + issued_at=issued_at, ) + self._sessions[self._digest(raw_token)] = principal + return SessionExchange(token=raw_token, principal=principal) def authenticate(self, token: str) -> SessionPrincipal: + """Validate a bearer token and slide its expiry forward. + + A session in continuous use must never expire mid-session -- every + successful authenticate() extends expires_at by the full TTL again, + capped at issued_at + max_lifetime so a session cannot renew itself + forever. A session that goes genuinely idle for longer than the TTL + still expires normally, since nothing calls authenticate() to renew + it while idle. + """ now = datetime.now(timezone.utc) + digest = self._digest(token) with self._lock: - principal = self._sessions.get(self._digest(token)) + principal = self._sessions.get(digest) if principal is None or principal.expires_at <= now: raise SessionSecurityError("invalid or expired session") + renewed_expiry = min(now + self._ttl, principal.issued_at + self._max_lifetime) + if renewed_expiry > principal.expires_at: + principal = replace(principal, expires_at=renewed_expiry) + self._sessions[digest] = principal return principal def validate_request_context(self, request: Request) -> None: diff --git a/backend/cortex_backend/core/paths.py b/backend/cortex_backend/core/paths.py index 25e8935..87a989e 100644 --- a/backend/cortex_backend/core/paths.py +++ b/backend/cortex_backend/core/paths.py @@ -79,6 +79,17 @@ def permanent_memory(self) -> Path: def permanent_memory_backup(self) -> Path: return self.data_dir / "memory_bank.json.bak" + @property + def settings_database(self) -> Path: + """Settings kept out of the chat database. + + Settings writes take a full-file backup copy first. Colocating them + with chat history meant every settings save byte-copied the entire + transcript store -- slow, disk-doubling, and able to fail a theme + toggle outright once the chat database grew large. + """ + return self.data_dir / "cortex_settings.sqlite" + @property def vector_database(self) -> Path: """Retain the dormant legacy path without enabling vector memory.""" diff --git a/backend/cortex_backend/execution/code_execution.py b/backend/cortex_backend/execution/code_execution.py index a70c883..8c69751 100644 --- a/backend/cortex_backend/execution/code_execution.py +++ b/backend/cortex_backend/execution/code_execution.py @@ -32,6 +32,7 @@ from urllib.error import HTTPError, URLError from urllib.parse import urljoin, urlsplit from urllib.request import ( + AbstractHTTPHandler, HTTPRedirectHandler, Request as UrlRequest, build_opener, @@ -999,7 +1000,16 @@ def _terminate_brokered_process(process: subprocess.Popen[bytes], job: _WindowsP pass -def _validate_network_url(url: str) -> str: +def _validate_network_url(url: str) -> tuple[str, str]: + """Validate a URL and return it with the single vetted IP to dial. + + Returning the resolved address is the point: the caller must connect to + *this* address rather than let the stack re-resolve the hostname, or a + time-varying DNS answer can pass the public-address check below and then + steer the actual connection to loopback/LAN (DNS rebinding). The + filesystem capability already closes the equivalent race by re-stat'ing + and comparing an identity tuple after validating. + """ if not isinstance(url, str) or len(url) > MAX_CODE_NETWORK_URL_CHARS: raise ValueError("network URL is invalid") try: @@ -1050,12 +1060,88 @@ def _validate_network_url(url: str) -> str: or resolved.is_unspecified ): raise PermissionError("network host is not public") - return url + # Every address in this answer passed, so any of them is safe to use. + # Pin the first so the connection cannot resolve a different one. + return url, addresses[0][4][0] + + +def _pinned_connection_classes(pinned_ip: str) -> tuple[type, type]: + """HTTP/HTTPS connection classes that dial ``pinned_ip`` directly. + + The hostname still travels in the ``Host`` header and in TLS SNI and + certificate validation, so servers and certificate checks behave exactly + as they normally would -- only the address the socket connects to is + forced, which is what closes the rebinding window. + """ + import http.client + + class _PinnedHTTPConnection(http.client.HTTPConnection): + def connect(self) -> None: + self.sock = self._create_connection( + (pinned_ip, self.port), self.timeout, self.source_address + ) + try: + self.sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) + except OSError: + pass + + class _PinnedHTTPSConnection(http.client.HTTPSConnection): + def connect(self) -> None: + self.sock = self._create_connection( + (pinned_ip, self.port), self.timeout, self.source_address + ) + try: + self.sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) + except OSError: + pass + # server_hostname stays the real hostname: certificate validation + # must not be weakened just because we dialed an address. + self.sock = self._context.wrap_socket(self.sock, server_hostname=self.host) + + return _PinnedHTTPConnection, _PinnedHTTPSConnection + + +class _PinnedHTTPHandler(AbstractHTTPHandler): + """Opens http:// requests through a caller-supplied connection factory.""" + + def __init__(self, connection_factory: Callable[..., Any]) -> None: + super().__init__() + self._connection_factory = connection_factory + + def http_open(self, req: Any) -> Any: + return self.do_open(self._connection_factory, req) + + http_request = AbstractHTTPHandler.do_request_ + + +class _PinnedHTTPSHandler(AbstractHTTPHandler): + """Opens https:// requests through a caller-supplied connection factory.""" + + def __init__(self, connection_factory: Callable[..., Any]) -> None: + super().__init__() + self._connection_factory = connection_factory + + def https_open(self, req: Any) -> Any: + return self.do_open(self._connection_factory, req) + + https_request = AbstractHTTPHandler.do_request_ class _SafeRedirectHandler(HTTPRedirectHandler): + """Re-validates every redirect target and re-pins the opener to it. + + Validating without re-pinning would leave the same hole one hop later: + the redirect's own connection would re-resolve the new hostname and could + land on an address this check just rejected. + """ + + def __init__(self, rebind: Callable[[str], None]) -> None: + super().__init__() + self._rebind = rebind + def redirect_request(self, request: Any, fp: Any, code: int, msg: str, headers: Any, newurl: str) -> Any: - _validate_network_url(urljoin(request.full_url, newurl)) + _, pinned_ip = _validate_network_url(urljoin(request.full_url, newurl)) + self._rebind(pinned_ip) return super().redirect_request(request, fp, code, msg, headers, newurl) @@ -1075,10 +1161,32 @@ def get(self, url: str, timeout: float = 5.0) -> str: raise ValueError("network timeout is invalid") from None if not math.isfinite(timeout): raise ValueError("network timeout is invalid") - safe_url = _validate_network_url(url) + safe_url, pinned_ip = _validate_network_url(url) self._budget.take_network() timeout = max(0.1, min(timeout, 5.0)) - opener = build_opener(_SafeRedirectHandler, ProxyHandler({})) + + # A mutable holder so a redirect can re-pin the opener to whatever its + # own (re-validated) target resolved to. The connection classes are + # rebuilt per connection rather than once up front -- otherwise every + # hop would keep dialing the first hop's address. + current_ip = {"value": pinned_ip} + + def rebind(next_ip: str) -> None: + current_ip["value"] = next_ip + + def connection_factory(secure: bool) -> Callable[..., Any]: + def factory(*args: Any, **kwargs: Any) -> Any: + plain, tls = _pinned_connection_classes(current_ip["value"]) + return (tls if secure else plain)(*args, **kwargs) + + return factory + + opener = build_opener( + _PinnedHTTPHandler(connection_factory(secure=False)), + _PinnedHTTPSHandler(connection_factory(secure=True)), + _SafeRedirectHandler(rebind), + ProxyHandler({}), + ) request = UrlRequest(safe_url, headers={"User-Agent": "Cortex-local-code/1"}) try: with opener.open(request, timeout=timeout) as response: diff --git a/backend/cortex_backend/launcher/frontend.py b/backend/cortex_backend/launcher/frontend.py index ae3d768..be53a5a 100644 --- a/backend/cortex_backend/launcher/frontend.py +++ b/backend/cortex_backend/launcher/frontend.py @@ -6,6 +6,7 @@ from datetime import datetime, timezone import hashlib import json +import logging import os from pathlib import Path import shutil @@ -14,8 +15,11 @@ from typing import Any +logger = logging.getLogger(__name__) + MANIFEST_NAME = ".cortex-build.json" INSTALL_MANIFEST_NAME = ".cortex-install.json" +INSTALL_CACHE_DIRNAME = ".cortex-install-cache" TRACKED_CONFIG = ( "index.html", "package.json", @@ -149,9 +153,29 @@ def _run(command: list[str], *, cwd: Path) -> None: ) from exc +def _reclaim_stale_staging_directories(parent: Path, keep: Path) -> None: + """Remove staging directories orphaned by a crashed or killed build.""" + try: + candidates = list(parent.glob(".cortex-frontend-build-*")) + except OSError as exc: + logger.warning("Could not scan %s for stale frontend build directories: %s", parent, exc) + return + for candidate in candidates: + if candidate == keep: + continue + try: + shutil.rmtree(candidate) + logger.info("Reclaimed stale frontend build staging directory: %s", candidate) + except OSError as exc: + logger.warning( + "Could not remove stale frontend build directory %s: %s", candidate, exc + ) + + def _stage_frontend_source(frontend_root: Path) -> Path: """Copy build inputs beside the source tree so live installs stay untouched.""" staging = frontend_root.parent / f".cortex-frontend-build-{uuid.uuid4().hex}" + _reclaim_stale_staging_directories(frontend_root.parent, staging) try: shutil.copytree( frontend_root, @@ -173,19 +197,30 @@ def _stage_frontend_source(frontend_root: Path) -> Path: return staging -def _install_if_needed(frontend_root: Path, expected_lock_digest: str) -> None: - node_modules = frontend_root / "node_modules" - marker = node_modules / INSTALL_MANIFEST_NAME - installed_digest = None +def _install_cache_root(frontend_root: Path) -> Path: + """Stable cache directory that survives per-build staging churn.""" + return frontend_root / INSTALL_CACHE_DIRNAME + + +def _install_if_needed(build_root: Path, expected_lock_digest: str, cache_root: Path) -> None: + """Install dependencies into ``build_root``, reusing a stable cache when possible.""" + node_modules = build_root / "node_modules" + cached_node_modules = cache_root / "node_modules" + marker = cache_root / INSTALL_MANIFEST_NAME + cached_digest = None if marker.is_file(): try: - installed_digest = json.loads(marker.read_text(encoding="utf-8"))["lock_digest"] + cached_digest = json.loads(marker.read_text(encoding="utf-8"))["lock_digest"] except (OSError, KeyError, TypeError, ValueError, json.JSONDecodeError): - installed_digest = None - if node_modules.is_dir() and installed_digest == expected_lock_digest: + cached_digest = None + if cached_digest == expected_lock_digest and cached_node_modules.is_dir(): + shutil.copytree(cached_node_modules, node_modules) return - _run([_tool_name("npm"), "ci"], cwd=frontend_root) - node_modules.mkdir(parents=True, exist_ok=True) + _run([_tool_name("npm"), "ci"], cwd=build_root) + cache_root.mkdir(parents=True, exist_ok=True) + if cached_node_modules.exists(): + shutil.rmtree(cached_node_modules) + shutil.copytree(node_modules, cached_node_modules) marker.write_text( json.dumps({"lock_digest": expected_lock_digest}, indent=2), encoding="utf-8", @@ -209,7 +244,7 @@ def build_frontend( source = source_digest(build_root) node_major = _major_version("node") npm_major = _major_version("npm") - _install_if_needed(build_root, lock) + _install_if_needed(build_root, lock, _install_cache_root(frontend_root)) _run( [_tool_name("npm"), "run", "build", "--", "--outDir", str(staging)], cwd=build_root, diff --git a/backend/cortex_backend/llamacpp/binary_fetcher.py b/backend/cortex_backend/llamacpp/binary_fetcher.py index 76736df..1e3cb15 100644 --- a/backend/cortex_backend/llamacpp/binary_fetcher.py +++ b/backend/cortex_backend/llamacpp/binary_fetcher.py @@ -62,28 +62,24 @@ def hash_directory(root: Path) -> str: return digest.hexdigest() -def _verify_directory(target_dir: Path, asset: AssetSpec) -> bool: - """Re-verify the whole extracted directory against the pinned manifest hash. - - Run before every launch (not just once after download) so a corrupted or - tampered-with cached install is caught rather than trusted forever. This - isn't a single-file TOCTOU-resistant stat-hash-stat check (a multi-file - tree walk can't be made atomic that cheaply); it is still a large - improvement over trusting an unverified cache indefinitely. +_TreeIdentity = tuple[tuple[str, int, int], ...] + + +def _tree_identity(root: Path) -> _TreeIdentity: + """Cheap per-file ``(relative_path, size, mtime_ns)`` fingerprint of a tree. + + Stat-ing every file is orders of magnitude cheaper than hashing their + content, so ``_verify_directory`` uses this to detect "nothing changed" + and skip the expensive ``hash_directory`` walk -- mirroring the + stat-based memoization ``GGUFModelDirectory`` uses for the same reason. """ - if not (target_dir / asset.executable_relpath).is_file(): - return False - try: - return hash_directory(target_dir) == asset.directory_sha256 - except (OSError, MemoryError): - # This check runs on every /api/v1/system poll (every 2s while a - # GGUF model is selected -- see App.tsx), including while a large - # local model is loaded and system memory is under real pressure. - # MemoryError is not an OSError subclass, so without this it was - # escaping uncaught and 500ing the whole system-status endpoint in - # a tight, permanent poll loop instead of just reporting "not - # verified as cached" the way a disk-read OSError already does. - return False + return tuple( + sorted( + (path.relative_to(root).as_posix(), (stat := path.stat()).st_size, stat.st_mtime_ns) + for path in root.rglob("*") + if path.is_file() + ) + ) class BinaryFetcher: @@ -92,17 +88,54 @@ class BinaryFetcher: def __init__(self, runtime_dir: Path, *, http_client: httpx.Client | None = None) -> None: self._runtime_dir = runtime_dir self._http = http_client + self._verification_cache: dict[Path, tuple[_TreeIdentity, bool]] = {} + + def _verify_directory(self, target_dir: Path, asset: AssetSpec) -> bool: + """Re-verify the whole extracted directory against the pinned manifest hash. + + Run before every launch (not just once after download) so a corrupted or + tampered-with cached install is caught rather than trusted forever. This + isn't a single-file TOCTOU-resistant stat-hash-stat check (a multi-file + tree walk can't be made atomic that cheaply); it is still a large + improvement over trusting an unverified cache indefinitely. + + The full SHA-256 walk (``hash_directory``) only actually runs when the + tree's cheap ``_tree_identity`` fingerprint has changed since the last + call -- this check runs on every /api/v1/system poll (every 2s while a + GGUF model is selected -- see App.tsx), and re-hashing a ~100MB, + unchanged runtime directory on every idle poll was pure wasted CPU + and disk I/O. + """ + if not (target_dir / asset.executable_relpath).is_file(): + return False + try: + identity = _tree_identity(target_dir) + cached = self._verification_cache.get(target_dir) + if cached is not None and cached[0] == identity: + return cached[1] + result = hash_directory(target_dir) == asset.directory_sha256 + self._verification_cache[target_dir] = (identity, result) + return result + except (OSError, MemoryError): + # This check runs on every /api/v1/system poll (every 2s while a + # GGUF model is selected -- see App.tsx), including while a large + # local model is loaded and system memory is under real pressure. + # MemoryError is not an OSError subclass, so without this it was + # escaping uncaught and 500ing the whole system-status endpoint in + # a tight, permanent poll loop instead of just reporting "not + # verified as cached" the way a disk-read OSError already does. + return False def is_cached(self, release: PinnedRelease, backend: GpuBackend) -> bool: asset = release.assets[backend] - return _verify_directory(self._target_dir(release, backend), asset) + return self._verify_directory(self._target_dir(release, backend), asset) def ensure_binary(self, release: PinnedRelease, backend: GpuBackend) -> Path: """Return a verified llama-server.exe path, downloading on first use.""" asset = release.assets[backend] target_dir = self._target_dir(release, backend) exe_path = target_dir / asset.executable_relpath - if _verify_directory(target_dir, asset): + if self._verify_directory(target_dir, asset): return exe_path self._runtime_dir.mkdir(parents=True, exist_ok=True) @@ -120,7 +153,7 @@ def ensure_binary(self, release: PinnedRelease, backend: GpuBackend) -> Path: if extract_tmp.exists(): shutil.rmtree(extract_tmp, ignore_errors=True) - if not _verify_directory(target_dir, asset): + if not self._verify_directory(target_dir, asset): raise BinaryVerificationError( f"Downloaded llama.cpp binary for '{backend}' failed verification." ) diff --git a/backend/cortex_backend/llamacpp/chat_client.py b/backend/cortex_backend/llamacpp/chat_client.py index fbb28ea..ef74551 100644 --- a/backend/cortex_backend/llamacpp/chat_client.py +++ b/backend/cortex_backend/llamacpp/chat_client.py @@ -10,9 +10,11 @@ from __future__ import annotations +import json import time from collections.abc import Callable from pathlib import Path +from threading import Event from typing import Any import httpx @@ -48,7 +50,14 @@ def set_status_callback(self, callback: Callable[[str], None] | None) -> None: """ self._status_callback = callback - def chat(self, *, model: str, messages: list[dict], options: dict) -> dict: + def chat( + self, + *, + model: str, + messages: list[dict], + options: dict, + cancellation_event: Event | None = None, + ) -> dict: model_path = resolve_gguf_path(self._models_directory(), model) # None means "no preference" -- title/translation calls pass a # minimal options dict with no num_ctx at all. Since num_ctx is a @@ -61,21 +70,21 @@ def chat(self, *, model: str, messages: list[dict], options: dict) -> dict: raw_num_ctx = options.get("num_ctx") num_ctx = int(raw_num_ctx) if raw_num_ctx is not None else None handle = self._provider.ensure_ready(model_path, num_ctx=num_ctx, on_status=self._status_callback) - body = _build_request_body(messages, options) started = time.monotonic() + if cancellation_event is None: + return self._chat_blocking(handle.base_url, messages, options, started) + return self._chat_abortable(handle.base_url, messages, options, started, cancellation_event) + + def _chat_blocking(self, base_url: str, messages: list[dict], options: dict, started: float) -> dict: + """Single request/response call, unchanged from before cancellation + support existed. Used whenever the caller has no cancellation_event + to honor (title and translation calls, and anything else that isn't + the main chat turn).""" + body = _build_request_body(messages, options, stream=False) try: - response = self._http.post(f"{handle.base_url}/v1/chat/completions", json=body) + response = self._http.post(f"{base_url}/v1/chat/completions", json=body) response.raise_for_status() except httpx.HTTPStatusError as exc: - # Carry llama-server's own explanation through instead of replacing - # it with a fixed string. Without it every failure -- a context - # overflow, an out-of-memory abort, an unsupported quantization -- - # arrived at _generation_failure_message() as the same opaque text, - # so none of its classifiers could match and every one of them was - # reported to the user as "rejected this request", which reads as - # if their message had been refused. The raw text is used only for - # classification there; the message shown to the user is always one - # of that function's curated strings, and this text is never logged. raise LlamaCppError( _server_error_detail(exc.response), status_code=exc.response.status_code, @@ -86,6 +95,77 @@ def chat(self, *, model: str, messages: list[dict], options: dict) -> dict: ) from exc return _adapt_to_ollama_shape(response.json(), elapsed_seconds=time.monotonic() - started) + def _chat_abortable( + self, + base_url: str, + messages: list[dict], + options: dict, + started: float, + cancellation_event: Event, + ) -> dict: + """Streamed request whose consumption is checked against + cancellation_event between chunks, so closing the response (which + releases llama-server's slot) happens promptly on Stop instead of + only after the model finishes on its own.""" + body = _build_request_body(messages, options, stream=True) + content_parts: list[str] = [] + reasoning_parts: list[str] = [] + usage: dict | None = None + timings: dict | None = None + try: + with self._http.stream("POST", f"{base_url}/v1/chat/completions", json=body) as response: + try: + response.raise_for_status() + except httpx.HTTPStatusError as exc: + exc.response.read() + raise LlamaCppError( + _server_error_detail(exc.response), + status_code=exc.response.status_code, + ) from exc + for line in response.iter_lines(): + if cancellation_event.is_set(): + break + if not line or not line.startswith("data:"): + continue + payload = line[len("data:"):].strip() + if payload == "[DONE]": + break + try: + chunk = json.loads(payload) + except ValueError: + continue + choices = chunk.get("choices") or [] + if choices: + delta = choices[0].get("delta") or {} + content_piece = delta.get("content") + if content_piece: + content_parts.append(content_piece) + reasoning_piece = delta.get("reasoning_content") + if reasoning_piece: + reasoning_parts.append(reasoning_piece) + if chunk.get("usage"): + usage = chunk["usage"] + if chunk.get("timings"): + timings = chunk["timings"] + except httpx.TransportError as exc: + raise LlamaCppError( + "Cortex lost its connection to the local model runtime." + ) from exc + # Reuse the existing non-streamed adapter by handing it a payload + # shaped the same way -- accumulated deltas standing in for the + # single message a non-streamed response would have carried. + synthetic_payload = { + "choices": [{ + "message": { + "content": "".join(content_parts), + "reasoning_content": "".join(reasoning_parts) or None, + }, + }], + "usage": usage, + "timings": timings, + } + return _adapt_to_ollama_shape(synthetic_payload, elapsed_seconds=time.monotonic() - started) + _MAX_SERVER_ERROR_CHARS = 400 @@ -115,11 +195,19 @@ def _server_error_detail(response: httpx.Response) -> str: return detail.strip()[:_MAX_SERVER_ERROR_CHARS] -def _build_request_body(messages: list[dict], options: dict) -> dict[str, Any]: +def _build_request_body(messages: list[dict], options: dict, *, stream: bool) -> dict[str, Any]: body: dict[str, Any] = { "messages": _strip_unsupported_fields(messages), - "stream": False, + "stream": stream, } + if stream: + # OpenAI-compatible streaming convention llama-server also follows: + # without this, per-chunk usage/timings are commonly omitted + # entirely rather than attached to the final chunk. Parsing already + # treats both as optional and falls back to a wall-clock estimate + # (see _adapt_to_ollama_shape), so an older server that ignores this + # field degrades to that same fallback rather than failing. + body["stream_options"] = {"include_usage": True} for option_key, body_key in ( ("temperature", "temperature"), ("top_p", "top_p"), diff --git a/backend/cortex_backend/llamacpp/server_manager.py b/backend/cortex_backend/llamacpp/server_manager.py index 2693ce5..03d8cd4 100644 --- a/backend/cortex_backend/llamacpp/server_manager.py +++ b/backend/cortex_backend/llamacpp/server_manager.py @@ -21,6 +21,7 @@ from __future__ import annotations +import ctypes import json import logging import socket @@ -28,9 +29,10 @@ import sys import threading import time +from ctypes import wintypes from dataclasses import dataclass from pathlib import Path -from typing import Callable, Literal, Protocol +from typing import Any, Callable, Literal, Protocol import httpx @@ -67,6 +69,12 @@ # already-loaded context size to inherit. In normal use the main chat call # establishes the real context size first, so this rarely matters. _DEFAULT_NUM_CTX = 4096 +# How long a vulkan launch failure for one (model, num_ctx, release) keeps +# steering that exact configuration to cpu before being retried on vulkan +# again -- long enough that a genuinely-too-large model doesn't thrash on +# every message, short enough that a driver update or freed VRAM gets a +# chance to matter within the same day rather than needing a manual reset. +_KNOWN_BAD_BACKEND_TTL_SECONDS = 24.0 * 3600.0 @dataclass(frozen=True, slots=True) @@ -131,7 +139,7 @@ def __call__(self, argv: list[str], *, cwd: Path) -> subprocess.Popen: ... -def default_launcher(argv: list[str], *, cwd: Path) -> subprocess.Popen: +def _spawn_process(argv: list[str], *, cwd: Path) -> subprocess.Popen: creationflags = subprocess.CREATE_NO_WINDOW if sys.platform == "win32" else 0 return subprocess.Popen( argv, @@ -142,6 +150,147 @@ def default_launcher(argv: list[str], *, cwd: Path) -> subprocess.Popen: ) +# Windows Job Object plumbing so llama-server cannot outlive this process. +# Sandboxed execution workers already get this exact policy (see +# execution/native_win32.py's Win32SuspendedWorker); it was simply missing +# here. Kept self-contained rather than importing that module's structs -- +# this manager deliberately does not depend on the AppContainer sandbox +# machinery (see the module docstring), and the struct layout below is +# stable, documented Win32 API surface, not something specific to either +# module. +_JOBOBJECT_EXTENDED_LIMIT_INFORMATION_CLASS = 9 +_JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE = 0x00002000 +_PROCESS_SET_QUOTA = 0x0100 +_PROCESS_TERMINATE = 0x0001 + + +class _JobObjectBasicLimitInformation(ctypes.Structure): + _fields_ = [ + ("per_process_user_time", ctypes.c_int64), + ("per_job_user_time", ctypes.c_int64), + ("limit_flags", wintypes.DWORD), + ("minimum_working_set_size", ctypes.c_size_t), + ("maximum_working_set_size", ctypes.c_size_t), + ("active_process_limit", wintypes.DWORD), + ("affinity", ctypes.c_size_t), + ("priority_class", wintypes.DWORD), + ("scheduling_class", wintypes.DWORD), + ] + + +class _JobObjectIoCounters(ctypes.Structure): + _fields_ = [ + ("read_operation_count", ctypes.c_uint64), + ("write_operation_count", ctypes.c_uint64), + ("other_operation_count", ctypes.c_uint64), + ("read_transfer_count", ctypes.c_uint64), + ("write_transfer_count", ctypes.c_uint64), + ("other_transfer_count", ctypes.c_uint64), + ] + + +class _JobObjectExtendedLimitInformation(ctypes.Structure): + _fields_ = [ + ("basic_limit_information", _JobObjectBasicLimitInformation), + ("io_info", _JobObjectIoCounters), + ("process_memory_limit", ctypes.c_size_t), + ("job_memory_limit", ctypes.c_size_t), + ("peak_process_memory_used", ctypes.c_size_t), + ("peak_job_memory_used", ctypes.c_size_t), + ] + + +class _JobWin32(Protocol): + """The handful of kernel32 entry points needed to assign a kill-on-close + Job Object -- small and injectable so tests can verify the exact call + sequence without touching real Windows APIs or spawning a real process.""" + + def CreateJobObjectW(self, security_attributes: Any, name: Any) -> int: ... + def SetInformationJobObject(self, job: int, info_class: int, info: Any, info_size: int) -> int: ... + def OpenProcess(self, access: int, inherit_handle: int, pid: int) -> int: ... + def AssignProcessToJobObject(self, job: int, process: int) -> int: ... + def CloseHandle(self, handle: int) -> int: ... + + +def _real_job_win32() -> _JobWin32: + kernel32 = ctypes.WinDLL("kernel32", use_last_error=True) + kernel32.CreateJobObjectW.argtypes = [wintypes.LPVOID, wintypes.LPCWSTR] + kernel32.CreateJobObjectW.restype = wintypes.HANDLE + kernel32.SetInformationJobObject.argtypes = [wintypes.HANDLE, ctypes.c_int, wintypes.LPVOID, wintypes.DWORD] + kernel32.SetInformationJobObject.restype = wintypes.BOOL + kernel32.OpenProcess.argtypes = [wintypes.DWORD, wintypes.BOOL, wintypes.DWORD] + kernel32.OpenProcess.restype = wintypes.HANDLE + kernel32.AssignProcessToJobObject.argtypes = [wintypes.HANDLE, wintypes.HANDLE] + kernel32.AssignProcessToJobObject.restype = wintypes.BOOL + kernel32.CloseHandle.argtypes = [wintypes.HANDLE] + kernel32.CloseHandle.restype = wintypes.BOOL + return kernel32 + + +class _JobObjectLauncher: + """Spawns llama-server and assigns it to a kill-on-close Job Object. + + The job is created once and held for the launcher's lifetime (one + instance per LlamaServerManager, held as a module-level default so a + normal Cortex process shares a single job across every model + restart) rather than recreated per launch, so restarting the server + many times in one session cannot leak a Windows handle per restart. + A hard exit of this Cortex process -- Task Manager, a crash, the + launcher supervisor's own shutdown timeout -- closes every handle this + process owns, including the job's; that is what tears llama-server + down with it even when nothing here ran a graceful stop() first. + """ + + def __init__(self, *, win32_factory: Callable[[], _JobWin32] = _real_job_win32) -> None: + self._win32_factory = win32_factory + self._win32: _JobWin32 | None = None + self._job: int | None = None + + def __call__(self, argv: list[str], *, cwd: Path) -> subprocess.Popen: + process = _spawn_process(argv, cwd=cwd) + if sys.platform == "win32": + self._apply_job_policy(process) + return process + + def _apply_job_policy(self, process: subprocess.Popen) -> None: + try: + win32 = self._win32 or self._win32_factory() + job = self._job + if job is None: + job = win32.CreateJobObjectW(None, None) + if not job: + return + limits = _JobObjectExtendedLimitInformation() + limits.basic_limit_information.limit_flags = _JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE + if not win32.SetInformationJobObject( + job, + _JOBOBJECT_EXTENDED_LIMIT_INFORMATION_CLASS, + ctypes.byref(limits), + ctypes.sizeof(limits), + ): + win32.CloseHandle(job) + return + self._win32 = win32 + self._job = job + process_handle = win32.OpenProcess( + _PROCESS_SET_QUOTA | _PROCESS_TERMINATE, False, process.pid + ) + if not process_handle: + return + try: + win32.AssignProcessToJobObject(job, process_handle) + finally: + win32.CloseHandle(process_handle) + except OSError: + logger.warning( + "Could not attach the local model runtime to a Job Object; " + "it may keep running if Cortex exits abnormally." + ) + + +default_launcher: ProcessLauncher = _JobObjectLauncher() + + def _free_loopback_port() -> int: with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: sock.bind(("127.0.0.1", 0)) @@ -412,22 +561,34 @@ def _guard_against_crash_loop(self, model_path: Path, effective_num_ctx: int) -> self._failure_times = [ at for at in self._failure_times if now - at < _FAILURE_WINDOW_SECONDS ] - if ( + tripped = ( self._failure_key == (model_path, effective_num_ctx) and len(self._failure_times) >= _FAILURE_LIMIT - ): - reason = self._last_restart_reason or "the runtime kept failing" - message = ( - f"The local model runtime for {model_path.name} failed " - f"{len(self._failure_times)} times in the last few minutes " - f"(most recently: {reason}). It likely does not fit in available " - "memory. Choose a smaller model or quantization, or lower the " - "context window in Settings, and Cortex will try again." - ) - self._terminate_and_reset_locked() - self._state = "failed" - self._last_error = message - raise LlamaCppError(message) + ) + if not tripped: + return + reason = self._last_restart_reason or "the runtime kept failing" + message = ( + f"The local model runtime for {model_path.name} failed " + f"{len(self._failure_times)} times in the last few minutes " + f"(most recently: {reason}). It likely does not fit in available " + "memory. Choose a smaller model or quantization, or lower the " + "context window in Settings, and Cortex will try again." + ) + process = self._process + self._state = "stopping" + # Terminate outside the state lock, same as _terminate_and_reset: the + # grace wait can take seconds, and status -- polled every couple of + # seconds by the UI -- must stay responsive throughout, not queue + # behind a shutdown the class documents this lock as never holding + # for more than microseconds. + if process is not None: + self._terminate_process(process) + with self._state_lock: + self._reset_fields_locked() + self._state = "failed" + self._last_error = message + raise LlamaCppError(message) def _terminate_and_reset(self) -> None: with self._state_lock: @@ -441,11 +602,6 @@ def _terminate_and_reset(self) -> None: with self._state_lock: self._reset_fields_locked() - def _terminate_and_reset_locked(self) -> None: - if self._process is not None: - self._terminate_process(self._process) - self._reset_fields_locked() - @staticmethod def _terminate_process(process: subprocess.Popen) -> None: process.terminate() @@ -478,7 +634,7 @@ def _start( requested_backend = self._gpu_backend_setting() last_exc: Exception | None = None - for backend in self._backend_order(requested_backend): + for backend in self._backend_order(requested_backend, model_path, num_ctx): try: return self._start_with_backend(model_path, num_ctx, backend, on_status) except ServerLaunchError as exc: @@ -495,27 +651,54 @@ def _start( self._last_error = message raise last_exc or LlamaCppError(message) - def _backend_order(self, requested: GpuBackendSetting) -> list[GpuBackend]: + def _backend_order( + self, requested: GpuBackendSetting, model_path: Path, num_ctx: int + ) -> list[GpuBackend]: if requested == "cpu": return ["cpu"] if requested == "vulkan": return ["vulkan"] - if self._known_bad_backend() == "vulkan": + if self._known_bad_backend(model_path, num_ctx) == "vulkan": return ["cpu"] return ["vulkan", "cpu"] - def _known_bad_backend(self) -> str | None: + def _known_bad_backend(self, model_path: Path, num_ctx: int) -> str | None: + """Only skip vulkan when THIS (model, context size, runtime build) + is the one that failed, and only for a bounded window -- a launch + failure for one oversized model must not permanently disable GPU + inference for every other model, and a driver update or freed VRAM + deserves a retry rather than an indefinite ban.""" try: data = json.loads(self._preferred_backend_file.read_text("utf-8")) except (OSError, ValueError): return None - return data.get("known_bad") if isinstance(data, dict) else None + if not isinstance(data, dict): + return None + if data.get("model") != str(model_path) or data.get("num_ctx") != num_ctx: + return None + if self._release is not None and data.get("release") != getattr(self._release, "tag", None): + return None + marked_at = data.get("at") + if ( + not isinstance(marked_at, (int, float)) + or isinstance(marked_at, bool) + or time.time() - marked_at > _KNOWN_BAD_BACKEND_TTL_SECONDS + ): + return None + return data.get("known_bad") - def _mark_backend_bad(self, backend: GpuBackend) -> None: + def _mark_backend_bad(self, backend: GpuBackend, model_path: Path, num_ctx: int) -> None: try: self._runtime_dir.mkdir(parents=True, exist_ok=True) self._preferred_backend_file.write_text( - json.dumps({"known_bad": backend}), encoding="utf-8" + json.dumps({ + "known_bad": backend, + "model": str(model_path), + "num_ctx": num_ctx, + "release": getattr(self._release, "tag", None) if self._release is not None else None, + "at": time.time(), + }), + encoding="utf-8", ) except OSError: logger.warning("Could not persist the known-bad GPU backend marker.") @@ -558,7 +741,7 @@ def _start_with_backend( exit_code = process.poll() if exit_code is not None: if backend == "vulkan": - self._mark_backend_bad("vulkan") + self._mark_backend_bad("vulkan", model_path, num_ctx) raise ServerLaunchError( "The local model runtime exited before it became ready.\n" + "\n".join(stderr_tail[-20:]) diff --git a/backend/cortex_backend/repositories/legacy_storage.py b/backend/cortex_backend/repositories/legacy_storage.py index 2283c3a..f9ea5ce 100644 --- a/backend/cortex_backend/repositories/legacy_storage.py +++ b/backend/cortex_backend/repositories/legacy_storage.py @@ -22,6 +22,7 @@ import math import struct import tempfile +import threading from datetime import datetime, timedelta, timezone from cortex_backend.core.paths import AppPaths @@ -898,6 +899,7 @@ def __init__( memory_file_path = str(resolved_paths.permanent_memory) self.memory_file_path = memory_file_path self.backup_file_path = f"{self.memory_file_path}.bak" + self._lock = threading.RLock() self.memos = self._load_memos() @staticmethod @@ -991,7 +993,8 @@ def get_memos(self) -> list[str]: Returns: A list of memo strings. """ - return list(self.memos) + with self._lock: + return list(self.memos) def add_memo(self, memo_text: str): """ @@ -1000,16 +1003,17 @@ def add_memo(self, memo_text: str): Args: memo_text (str): The fact to be remembered. """ - normalized = self.normalize_memos(self.memos + [memo_text]) - if normalized == self.memos: - return - previous_memos = list(self.memos) - self.memos = normalized - try: - self._save_memos() - except PersistenceError: - self.memos = previous_memos - raise + with self._lock: + normalized = self.normalize_memos(self.memos + [memo_text]) + if normalized == self.memos: + return + previous_memos = list(self.memos) + self.memos = normalized + try: + self._save_memos() + except PersistenceError: + self.memos = previous_memos + raise def update_memos(self, memos: list[str]): """ @@ -1019,24 +1023,26 @@ def update_memos(self, memos: list[str]): memos (list[str]): The new, complete list of memos. """ # Filter out any empty strings that might have come from the UI. - previous_memos = list(self.memos) - self.memos = self.normalize_memos(memos) - try: - self._save_memos() - except PersistenceError: - self.memos = previous_memos - raise - logging.info(f"Permanent memory updated with {len(self.memos)} memos.") + with self._lock: + previous_memos = list(self.memos) + self.memos = self.normalize_memos(memos) + try: + self._save_memos() + except PersistenceError: + self.memos = previous_memos + raise + logging.info(f"Permanent memory updated with {len(self.memos)} memos.") def clear_memos(self): """Clears all memos from the list and saves the empty list to disk.""" - previous_memos = list(self.memos) - self.memos.clear() - try: - self._save_memos() - except PersistenceError: - self.memos = previous_memos - raise + with self._lock: + previous_memos = list(self.memos) + self.memos.clear() + try: + self._save_memos() + except PersistenceError: + self.memos = previous_memos + raise class ShortTermMemory: diff --git a/backend/cortex_backend/repositories/sqlite_settings.py b/backend/cortex_backend/repositories/sqlite_settings.py index 6768dc5..797cf96 100644 --- a/backend/cortex_backend/repositories/sqlite_settings.py +++ b/backend/cortex_backend/repositories/sqlite_settings.py @@ -24,6 +24,7 @@ SETTINGS_SCHEMA_VERSION = 1 MIGRATION_KEY = "qsettings-to-sqlite-v1" +COLOCATED_MIGRATION_KEY = "chatdb-colocated-settings-to-own-file-v1" def _utc_now() -> str: @@ -31,10 +32,15 @@ def _utc_now() -> str: class SQLiteSettingsRepository: - """Store validated settings beside the existing chat database. + """Store validated settings in their own database file. The repository creates only additive settings tables. It never writes back to QSettings, so the legacy Qt reader remains a safe rollback path. + + Settings used to live inside the chat database. Every save takes a + full-file backup copy first, so colocation meant each settings write + byte-copied the whole transcript store. ``adopt_from`` performs the + one-time move; see :meth:`_adopt_colocated_settings`. """ def __init__( @@ -42,6 +48,7 @@ def __init__( db_path: str | Path, *, legacy: SettingsRepository | None = None, + adopt_from: str | Path | None = None, ) -> None: self.db_path = Path(db_path) self.backup_path = Path(f"{self.db_path}.bak") @@ -52,6 +59,90 @@ def __init__( self._load_lock = RLock() self._pre_schema_backup = self._create_backup() self._ensure_schema() + if adopt_from is not None: + self._adopt_colocated_settings(Path(adopt_from)) + + def _adopt_colocated_settings(self, source_db: Path) -> None: + """Move settings out of a database they used to share with chat data. + + Runs once per install: if this settings database has no row yet but + the old colocated database does, copy that row across. Without this, + every existing install would silently revert to defaults on upgrade, + which is a worse failure than the one being fixed. + + The source row is left in place. It is small, it costs nothing to + keep, and leaving it makes downgrading to a previous Cortex build a + non-event rather than a data-loss bug. + """ + if source_db == self.db_path or not source_db.exists(): + return + with self._load_lock: + try: + with self.connect() as connection: + already = connection.execute( + "SELECT 1 FROM cortex_settings WHERE id = 1" + ).fetchone() + if already is not None: + return + except SettingsRepositoryError: + return + + source: sqlite3.Connection | None = None + try: + source = sqlite3.connect(source_db, timeout=10.0) + source.row_factory = sqlite3.Row + source.execute("PRAGMA busy_timeout = 10000") + table = source.execute( + "SELECT name FROM sqlite_master WHERE type = 'table' AND name = 'cortex_settings'" + ).fetchone() + if table is None: + return + row = source.execute( + "SELECT schema_version, revision, payload, updated_at " + "FROM cortex_settings WHERE id = 1" + ).fetchone() + except sqlite3.Error: + # The old database being unreadable must not stop Cortex from + # starting -- it just means there is nothing to adopt, and the + # normal legacy/default path takes over. + return + finally: + if source is not None: + source.close() + + if row is None: + return + try: + with self.connect() as connection: + connection.execute( + "INSERT OR IGNORE INTO cortex_settings " + "(id, schema_version, revision, payload, updated_at) " + "VALUES (1, ?, ?, ?, ?)", + ( + int(row["schema_version"]), + int(row["revision"]), + str(row["payload"]), + str(row["updated_at"]), + ), + ) + connection.execute( + "INSERT OR REPLACE INTO settings_migration_ledger " + "(migration_key, source, status, imported_keys, invalid_keys, " + "backup_path, message, applied_at) " + "VALUES (?, ?, ?, ?, ?, ?, ?, ?)", + ( + COLOCATED_MIGRATION_KEY, + str(source_db), + "applied", + "[]", + "[]", + None, + "Adopted settings from the chat database.", + _utc_now(), + ), + ) + except SettingsRepositoryError: + return @contextmanager def connect(self) -> Iterator[sqlite3.Connection]: diff --git a/backend/cortex_backend/services/chat_client.py b/backend/cortex_backend/services/chat_client.py index 6d32841..2aadbe1 100644 --- a/backend/cortex_backend/services/chat_client.py +++ b/backend/cortex_backend/services/chat_client.py @@ -9,6 +9,7 @@ from __future__ import annotations +from threading import Event from typing import Any, Protocol # Ollama tags are ``name:tag`` and never contain this prefix, so it @@ -25,9 +26,22 @@ class ChatClient(Protocol): "prompt_eval_duration": int | None, # nanoseconds "eval_duration": int | None, # nanoseconds "total_duration": int | None} # nanoseconds + + ``cancellation_event``, when given, lets a caller ask the client to stop + consuming an in-flight response early (see ``LlamaCppChatClient`` and + ``OllamaChatClient``). It is optional and only meaningful to real + implementations -- callers that never set it keep today's simple + single-shot request. """ - def chat(self, *, model: str, messages: list[dict], options: dict) -> dict: + def chat( + self, + *, + model: str, + messages: list[dict], + options: dict, + cancellation_event: Event | None = None, + ) -> dict: ... @@ -37,8 +51,49 @@ class OllamaChatClient: def __init__(self, client: Any) -> None: self._client = client - def chat(self, *, model: str, messages: list[dict], options: dict) -> dict: - return self._client.chat(model=model, messages=messages, options=options) + def chat( + self, + *, + model: str, + messages: list[dict], + options: dict, + cancellation_event: Event | None = None, + ) -> dict: + if cancellation_event is None: + return self._client.chat(model=model, messages=messages, options=options) + # ollama.Client(stream=True) returns a generator that owns an httpx + # streaming response internally (see the installed ``ollama`` package's + # Client._request: ``with self._client.stream(...) as r: ... yield``). + # Breaking out of the loop early and closing the generator sends it a + # GeneratorExit at its suspended yield point, which unwinds that + # ``with`` block and releases the connection -- the same mechanism + # LlamaCppChatClient uses for the local runtime. + chunks = self._client.chat(model=model, messages=messages, options=options, stream=True) + content_parts: list[str] = [] + thinking_parts: list[str] = [] + final: dict = {} + try: + for chunk in chunks: + if cancellation_event.is_set(): + break + message = chunk.get("message") or {} + content_piece = message.get("content") + if content_piece: + content_parts.append(content_piece) + thinking_piece = message.get("thinking") + if thinking_piece: + thinking_parts.append(thinking_piece) + if chunk.get("done"): + final = dict(chunk) + finally: + close = getattr(chunks, "close", None) + if callable(close): + close() + final["message"] = { + "content": "".join(content_parts), + "thinking": "".join(thinking_parts) or None, + } + return final class RoutingChatClient: @@ -55,8 +110,20 @@ def __init__(self, ollama_client: ChatClient, llamacpp_client: ChatClient) -> No self._ollama = ollama_client self._llamacpp = llamacpp_client - def chat(self, *, model: str, messages: list[dict], options: dict) -> dict: + def chat( + self, + *, + model: str, + messages: list[dict], + options: dict, + cancellation_event: Event | None = None, + ) -> dict: target = self._llamacpp if model.startswith(GGUF_PREFIX) else self._ollama + # Only forward cancellation_event when it is actually set, so test + # doubles and any future ChatClient implementation that predates this + # parameter keep working against their original 3-argument call. + if cancellation_event is not None: + return target.chat(model=model, messages=messages, options=options, cancellation_event=cancellation_event) return target.chat(model=model, messages=messages, options=options) def set_status_callback(self, callback: Any) -> None: diff --git a/backend/cortex_backend/services/generation.py b/backend/cortex_backend/services/generation.py index 722d552..dc7438c 100644 --- a/backend/cortex_backend/services/generation.py +++ b/backend/cortex_backend/services/generation.py @@ -47,9 +47,25 @@ def fit_history_to_context( num_ctx: int, code_execution_eligible: bool | None = None, bypass_system_prompt: bool = False, + attachments: Sequence[GenerationAttachment] = (), ) -> str: """Format the retained history for the model prompt.""" + def fit_attachments_to_context( + self, + attachments: Sequence[GenerationAttachment], + *, + query: str, + chat_history: str, + permanent_memories: list[str], + memories_enabled: bool, + user_system_instructions: str | None, + num_ctx: int, + code_execution_eligible: bool | None = None, + bypass_system_prompt: bool = False, + ) -> tuple[GenerationAttachment, ...]: + """Bound attachment reference text to fit the configured context.""" + def generate( self, *, @@ -60,6 +76,7 @@ def generate( user_system_instructions: str | None, options: dict[str, Any], attachments: Sequence[GenerationAttachment] = (), + cancellation_event: Event | None = None, ) -> tuple[str, str | None, MemoryCommand, GenerationStats | None]: """Generate a response and validated memory command.""" @@ -151,16 +168,41 @@ def generate( working_history = [dict(message) for message in loaded_history] if working_history and working_history[-1].get("role") == "user": working_history.pop() - chat_history = engine.fit_history_to_context( - working_history, - query=snapshot.user_input, - permanent_memories=permanent_memories, - memories_enabled=snapshot.memories_enabled, - user_system_instructions=snapshot.user_system_instructions, - num_ctx=num_ctx, - code_execution_eligible=snapshot.code_execution_eligible, - bypass_system_prompt=snapshot.bypass_system_prompt, - ) + + # Reserve room for attachments *before* history claims the whole + # budget: fit them first against a placeholder (history is not known + # yet), giving an attached document priority over old chat turns, + # then let history size itself around that reservation below. The + # attachments passed to engine.generate() further down are re-fit + # against the real, now-correctly-sized chat_history -- this pass + # only determines how much room history should leave. + reserved_attachments: Sequence[GenerationAttachment] = () + fit_attachments = getattr(engine, "fit_attachments_to_context", None) + if snapshot.attachments and callable(fit_attachments): + reserved_attachments = fit_attachments( + snapshot.attachments, + query=snapshot.user_input, + chat_history="No history available.", + permanent_memories=permanent_memories, + memories_enabled=snapshot.memories_enabled, + user_system_instructions=snapshot.user_system_instructions, + num_ctx=num_ctx, + code_execution_eligible=snapshot.code_execution_eligible, + bypass_system_prompt=snapshot.bypass_system_prompt, + ) + + history_kwargs: dict[str, Any] = { + "query": snapshot.user_input, + "permanent_memories": permanent_memories, + "memories_enabled": snapshot.memories_enabled, + "user_system_instructions": snapshot.user_system_instructions, + "num_ctx": num_ctx, + "code_execution_eligible": snapshot.code_execution_eligible, + "bypass_system_prompt": snapshot.bypass_system_prompt, + } + if reserved_attachments: + history_kwargs["attachments"] = reserved_attachments + chat_history = engine.fit_history_to_context(working_history, **history_kwargs) self._check_cancelled(cancellation_event) generate_kwargs: dict[str, Any] = { @@ -172,9 +214,12 @@ def generate( "options": dict(snapshot.model_options), } # Keep the legacy headless engine protocol compatible for callers that - # do not use attachments; real engines receive the resolved payload. + # do not use attachments or cancellation; real engines receive the + # resolved payload. if snapshot.attachments: generate_kwargs["attachments"] = snapshot.attachments + if cancellation_event is not None: + generate_kwargs["cancellation_event"] = cancellation_event response, thoughts, memory_command, stats = engine.generate( **generate_kwargs, ) diff --git a/backend/cortex_backend/services/llm.py b/backend/cortex_backend/services/llm.py index 8c9901a..e053375 100644 --- a/backend/cortex_backend/services/llm.py +++ b/backend/cortex_backend/services/llm.py @@ -14,6 +14,8 @@ import re import sys from collections.abc import Sequence +from threading import Event +from typing import Any from cortex_backend.core.generation import ( CodeExecutionProposal, @@ -540,8 +542,16 @@ def fit_history_to_context( num_ctx: int, code_execution_eligible: bool | None = None, bypass_system_prompt: bool = False, + attachments: Sequence[GenerationAttachment] = (), ) -> str: - """Keep the newest history that fits beside prompts, memories, and output.""" + """Keep the newest history that fits beside prompts, memories, and output. + + ``attachments`` are already-fitted reference text (see + ``fit_attachments_to_context``); they are threaded into the same + per-candidate prompt sizing used here purely so history leaves them + room, mirroring the fixed overhead memories and the system prompt + already contribute. + """ output_reservation = cls.output_token_reservation(num_ctx) selected: list[dict] = [] @@ -554,14 +564,20 @@ def fit_history_to_context( permanent_memories, memories_enabled, user_system_instructions, + attachments, code_execution_eligible=code_execution_eligible, bypass_system_prompt=bypass_system_prompt, ) prompt_tokens = sum(cls.estimate_tokens(item.get("content", "")) + 4 for item in prompt) if prompt_tokens + output_reservation <= max(256, int(num_ctx)): selected = candidate - elif selected: - break + # Candidate sizes are not monotonic: dropping a newly-unpaired + # trailing assistant message (see _format_history_messages) + # shrinks the *next* candidate, so an oversized exchange must not + # stop the walk -- older, smaller exchanges further back can + # still fit. Stopping here previously discarded the entire + # history whenever the single newest exchange alone was too + # large for the budget. return cls._format_history_messages(selected) @@ -627,6 +643,7 @@ def generate( user_system_instructions: str | None, options: dict | None = None, attachments: Sequence[GenerationAttachment] = (), + cancellation_event: Event | None = None, ) -> tuple[str, str | None, MemoryCommand, GenerationStats | None]: """ Generates a synthesized response and extracts thoughts and commands. @@ -638,6 +655,11 @@ def generate( memories_enabled (bool): Flag indicating if memory features are active. user_system_instructions (str | None): Custom instructions from the user. options (dict | None): A dictionary of Ollama options (e.g., temperature, num_ctx). + cancellation_event (Event | None): When given, lets the underlying + chat client stop consuming an in-flight response early instead + of only noticing cancellation after the call returns on its + own. Only the real chat turn passes one; title and + translation calls do not need it. Returns: A tuple containing: @@ -692,11 +714,14 @@ def generate( # the model call still "succeeds", but the persisted message has # empty content next to a full reasoning trace. - response = self.chat_client.chat( - model=self.gen_model, - messages=prompt_messages, - options=api_options - ) + chat_kwargs: dict[str, Any] = { + "model": self.gen_model, + "messages": prompt_messages, + "options": api_options, + } + if cancellation_event is not None: + chat_kwargs["cancellation_event"] = cancellation_event + response = self.chat_client.chat(**chat_kwargs) message_obj = response.get('message', {}) main_content = message_obj.get('content', '') thinking_content = message_obj.get('thinking') diff --git a/backend/cortex_backend/testing/fake_llamacpp.py b/backend/cortex_backend/testing/fake_llamacpp.py index ffbc6e0..74fba70 100644 --- a/backend/cortex_backend/testing/fake_llamacpp.py +++ b/backend/cortex_backend/testing/fake_llamacpp.py @@ -6,11 +6,13 @@ from __future__ import annotations +import json from dataclasses import dataclass, field from pathlib import Path from typing import Any from fastapi import FastAPI, HTTPException +from fastapi.responses import StreamingResponse from cortex_backend.llamacpp.server_manager import LlamaCppRuntimeStatus, ServerHandle @@ -71,7 +73,7 @@ def health() -> dict[str, str]: return {"status": "ok"} @app.post("/v1/chat/completions", response_model=None) - def chat_completions(payload: dict[str, Any]) -> dict[str, Any]: + def chat_completions(payload: dict[str, Any]) -> dict[str, Any] | StreamingResponse: messages = payload.get("messages") if not isinstance(messages, list) or not messages: raise HTTPException(status_code=422, detail="messages required") @@ -82,20 +84,47 @@ def chat_completions(payload: dict[str, Any]) -> dict[str, Any]: "", ) content = fake_state.generation_response or f"Echo: {last_user}" + usage = {"prompt_tokens": 24, "completion_tokens": 48, "total_tokens": 72} + timings = { + "prompt_n": 24, + "prompt_ms": 120.0, + "prompt_per_second": 200.0, + "predicted_n": 48, + "predicted_ms": 480.0, + "predicted_per_second": 100.0, + } + if payload.get("stream"): + def sse_chunks(): + if fake_state.generation_thoughts: + yield _sse({"choices": [{"delta": {"reasoning_content": fake_state.generation_thoughts}}]}) + for piece in _fake_llamacpp_chunks(content): + yield _sse({"choices": [{"delta": {"content": piece}}]}) + yield _sse({ + "choices": [{"delta": {}, "finish_reason": "stop"}], + "usage": usage, + "timings": timings, + }) + yield "data: [DONE]\n\n" + return StreamingResponse(sse_chunks(), media_type="text/event-stream") + message: dict[str, Any] = {"role": "assistant", "content": content} if fake_state.generation_thoughts: message["reasoning_content"] = fake_state.generation_thoughts return { "choices": [{"message": message, "finish_reason": "stop"}], - "usage": {"prompt_tokens": 24, "completion_tokens": 48, "total_tokens": 72}, - "timings": { - "prompt_n": 24, - "prompt_ms": 120.0, - "prompt_per_second": 200.0, - "predicted_n": 48, - "predicted_ms": 480.0, - "predicted_per_second": 100.0, - }, + "usage": usage, + "timings": timings, } return app + + +def _sse(payload: dict[str, Any]) -> str: + return f"data: {json.dumps(payload)}\n\n" + + +def _fake_llamacpp_chunks(value: str, size: int = 4): + """Split into several pieces so a streaming test observes more than one + chunk, mirroring the real server's token-at-a-time delivery.""" + for start in range(0, len(value), size): + yield value[start:start + size] diff --git a/backend/cortex_backend/testing/fake_ollama.py b/backend/cortex_backend/testing/fake_ollama.py index c1083e1..f0c0303 100644 --- a/backend/cortex_backend/testing/fake_ollama.py +++ b/backend/cortex_backend/testing/fake_ollama.py @@ -178,6 +178,7 @@ def generate( user_system_instructions: str | None, options: dict[str, Any], attachments: tuple[Any, ...] = (), + cancellation_event: Any = None, ) -> tuple[str, str | None, MemoryCommand, GenerationStats | None]: del ( chat_history, @@ -186,6 +187,7 @@ def generate( user_system_instructions, options, attachments, + cancellation_event, ) if self.state.status_updates and self._status_callback is not None: for message in self.state.status_updates: diff --git a/frontend/src/features/chat/ChatPage.test.tsx b/frontend/src/features/chat/ChatPage.test.tsx index a74e1cd..0919d2a 100644 --- a/frontend/src/features/chat/ChatPage.test.tsx +++ b/frontend/src/features/chat/ChatPage.test.tsx @@ -5,6 +5,7 @@ import { useState } from "react"; import type { ChatAttachment, ChatResponse } from "../../../../contracts/cortex-api"; import { ApiError, CortexApi } from "../../api/client"; import { humanizeGenerationStatus } from "../../lib/generationStatus"; +import { NEW_THREAD_OPTIONS_KEY, useChatStore } from "../../stores/useChatStore"; import { ChatPage } from "./ChatPage"; describe("humanizeGenerationStatus", () => { @@ -63,7 +64,10 @@ function renderChat(api: CortexApi, threadId = "thread-a", selectedModelSupports } describe("ChatPage composer integration", () => { - afterEach(() => window.sessionStorage.clear()); + afterEach(() => { + window.sessionStorage.clear(); + useChatStore.setState({ generationOptionsByThread: {} }); + }); it("keeps a blank conversation focused on the composer", async () => { renderChat(chatApi()); @@ -334,6 +338,53 @@ describe("ChatPage composer integration", () => { expect(JSON.parse(window.sessionStorage.getItem("cortex.composer.attachments.thread-new") ?? "[]")).toEqual([lateAttachment]); }); + it("migrates generation overrides set before the first message to the new chat's thread id", async () => { + const user = userEvent.setup(); + const api = chatApi({ + chat: vi.fn(async (id: string) => emptyChat(id)), + generate: vi.fn(async () => ({ + job_id: "job-new", + kind: "generation" as const, + status: "queued" as const, + thread_id: "thread-new", + user_message_id: "message-new", + })), + }); + function RoutedChat() { + const [threadId, setThreadId] = useState(null); + return ( + true} + onRescanModels={async () => undefined} + onThreadCreated={setThreadId} + onChatChanged={vi.fn()} + onForked={vi.fn()} + onSessionExpired={vi.fn()} + /> + ); + } + render(); + + // Tune sampling before the thread exists -- it is stored under the + // "new chat" placeholder key until a real thread id is assigned. + useChatStore.getState().setThreadOptions(NEW_THREAD_OPTIONS_KEY, { temperature: 0.2 }); + + const composer = await screen.findByLabelText("Message Cortex"); + await user.type(composer, "First turn"); + await user.click(screen.getByRole("button", { name: "Send message" })); + await waitFor(() => expect(api.generate).toHaveBeenCalledTimes(1)); + + await waitFor(() => expect(useChatStore.getState().generationOptionsByThread["thread-new"]).toEqual({ temperature: 0.2 })); + expect(useChatStore.getState().generationOptionsByThread[NEW_THREAD_OPTIONS_KEY]).toBeUndefined(); + }); + it("retargets an in-flight new-chat attachment when acceptance wins the race", async () => { const user = userEvent.setup(); const stagedAttachment: ChatAttachment = { @@ -403,7 +454,16 @@ describe("ChatPage composer integration", () => { expect(JSON.parse(window.sessionStorage.getItem("cortex.composer.attachments.thread-inverse") ?? "[]")).toEqual([stagedAttachment]); }); - it("replays an active generation from the beginning after a remount", async () => { + it("replays from the beginning on a cold start, then resumes without duplicating after a route remount", async () => { + // Two different situations that both reach this mount effect: + // + // Cold start (page reload): sessionStorage remembers the job but the + // module-level store was wiped, so the transcript must be rebuilt by + // replaying every event from 0. + // + // Route remount (Settings and back): the page unmounts but the store + // survives with its accumulated text intact. Replaying from 0 here + // appends the whole answer onto itself -- the regression this pins. window.sessionStorage.setItem("cortex.active.generation", JSON.stringify({ jobId: "job-replay", threadId: "thread-a", lastEventId: 7 })); const streamCalls: Array<{ afterEventId?: number }> = []; const api = chatApi({ @@ -415,12 +475,19 @@ describe("ChatPage composer integration", () => { }); const first = renderChat(api, "thread-a"); await waitFor(() => expect(streamCalls).toHaveLength(1)); + await waitFor(() => expect(useChatStore.getState().generation.partialContent).toBe("replayed")); + first.unmount(); renderChat(api, "thread-a"); await waitFor(() => expect(streamCalls).toHaveLength(2)); + // Cold start replayed from 0; the remount resumed from the persisted + // cursor instead of rewinding. expect(streamCalls[0].afterEventId).toBe(0); - expect(streamCalls[1].afterEventId).toBe(0); + expect(streamCalls[1].afterEventId).toBe(8); + + // And the answer is not printed twice. + await waitFor(() => expect(useChatStore.getState().generation.partialContent).toBe("replayed")); }); it("collapses the pending reasoning panel in step with contentReady, ahead of the swap to the real message", async () => { diff --git a/frontend/src/features/chat/ChatPage.tsx b/frontend/src/features/chat/ChatPage.tsx index b07c3ef..87a64f0 100644 --- a/frontend/src/features/chat/ChatPage.tsx +++ b/frontend/src/features/chat/ChatPage.tsx @@ -237,9 +237,17 @@ export function ChatPage({ void loadChat({ preserveCurrent }); const stored = readActiveJob(); if (stored) { - const job: PersistedJob = initialMountRef.current ? { ...stored, lastEventId: 0 } : stored; + // Only rewind the event cursor when the store is actually cold for + // this job. initialMountRef is per-instance, so it is true on EVERY + // mount -- including a return from /settings, which unmounts this + // page but leaves the module-level generation store populated. + // Replaying from 0 onto already-accumulated text printed the answer + // twice; keeping the stored cursor resumes where the buffer left off. + const warmForThisJob = useChatStore.getState().generation.jobId === stored.jobId; + const replayFromStart = initialMountRef.current && !warmForThisJob; + const job: PersistedJob = replayFromStart ? { ...stored, lastEventId: 0 } : stored; initialMountRef.current = false; - if (useChatStore.getState().generation.jobId !== job.jobId) { + if (replayFromStart || !warmForThisJob) { useChatStore.getState().beginGeneration(job.jobId, job.threadId); } void consume(job, reconcileChat, reportGenerationFailure); @@ -335,6 +343,17 @@ export function ChatPage({ const destinationThreadId = submittedThreadId ?? started.threadId; const destinationDraftScope = composerDraftKey(destinationThreadId); const destinationAttachmentScope = composerAttachmentKey(destinationThreadId); + if (!submittedThreadId) { + // Overrides staged before the first message live under the "new chat" + // placeholder key. Migrate them to the real thread id now that one + // exists, so they keep applying to this conversation instead of + // reverting after one message and leaking into the next new chat. + const draftOptions = useChatStore.getState().generationOptionsByThread[NEW_THREAD_OPTIONS_KEY]; + if (draftOptions) { + setThreadOptions(destinationThreadId, draftOptions); + setThreadOptions(NEW_THREAD_OPTIONS_KEY, null); + } + } if (submittedAttachmentScope !== destinationAttachmentScope) { // Retarget only batches that were already staging into this submitted // draft. Each batch owns its mutable target, so a later /chat/new never diff --git a/frontend/src/features/chat/MessageList.test.tsx b/frontend/src/features/chat/MessageList.test.tsx index bb20ba1..bc4346c 100644 --- a/frontend/src/features/chat/MessageList.test.tsx +++ b/frontend/src/features/chat/MessageList.test.tsx @@ -1,6 +1,7 @@ -import { createRef } from "react"; +import { createRef, forwardRef, useEffect, useImperativeHandle } from "react"; import { render, screen, waitFor } from "@testing-library/react"; import { describe, expect, it, vi } from "vitest"; +import type { VirtuosoHandle, VirtuosoProps } from "react-virtuoso"; import type { ChatMessage } from "../../../../contracts/cortex-api"; import { MessageList, type MessageListHandle } from "./MessageList"; @@ -97,6 +98,78 @@ describe("MessageList", () => { await waitFor(() => expect(screen.getByTestId("pending-bubble")).toBeInTheDocument()); }); + it("updates the virtualized Footer content in place instead of remounting it on every change", async () => { + // The real react-virtuoso only redraws its Footer slot in response to + // its own internal layout/scroll signals, which jsdom's zero-height + // environment never fires after mount -- so a plain rerender() can't + // observe an update through the real library here. Swap in a minimal + // stand-in that always re-invokes components.Footer on every render, + // the way the real library does in a browser, so this test can isolate + // and verify the actual contract MessageList relies on: components.Footer + // must be read fresh each render (so content updates) while the + // function's own identity stays stable (so React doesn't remount it). + vi.resetModules(); + vi.doMock("react-virtuoso", () => ({ + Virtuoso: forwardRef>(function MockVirtuoso(props, ref) { + useImperativeHandle(ref, () => ({ + scrollToIndex: () => {}, + scrollTo: () => {}, + scrollBy: () => {}, + autoscrollToBottom: () => {}, + scrollIntoView: () => {}, + getState: () => { throw new Error("not implemented in this test double"); }, + })); + const Footer = props.components?.Footer; + return
{Footer ?
: null}
; + }), + })); + const { MessageList: MockedMessageList } = await import("./MessageList"); + + const mountSpy = vi.fn(); + function Probe({ label }: { label: string }) { + // Fires only on true mount (empty deps) -- a remount would call this + // again; an in-place update of the same instance would not. + useEffect(() => { mountSpy(); }, []); + return
{label}
; + } + + const { rerender } = render( + } + />, + ); + expect(screen.getByTestId("pending-bubble")).toHaveTextContent("Streaming…"); + expect(mountSpy).toHaveBeenCalledTimes(1); + + rerender( + } + />, + ); + + expect(screen.getByTestId("pending-bubble")).toHaveTextContent("Streaming… more text"); + expect(mountSpy).toHaveBeenCalledTimes(1); + + vi.doUnmock("react-virtuoso"); + vi.resetModules(); + }); + it("reports near-end scroll state via onNearEndChange on the plain path", () => { const onNearEndChange = vi.fn(); render( diff --git a/frontend/src/features/chat/MessageList.tsx b/frontend/src/features/chat/MessageList.tsx index 44d5e38..cd3e1bf 100644 --- a/frontend/src/features/chat/MessageList.tsx +++ b/frontend/src/features/chat/MessageList.tsx @@ -1,4 +1,4 @@ -import { forwardRef, useImperativeHandle, useRef, type ReactNode } from "react"; +import { forwardRef, useCallback, useImperativeHandle, useRef, type ReactNode } from "react"; import { Virtuoso, type VirtuosoHandle } from "react-virtuoso"; import type { ChatMessage } from "../../../../contracts/cortex-api"; import { MessageCard } from "./MessageCard"; @@ -38,6 +38,18 @@ export const MessageList = forwardRef(function Message const virtuosoRef = useRef(null); const virtualized = messages.length >= VIRTUALIZE_THRESHOLD; + // Virtuoso remounts its Footer subtree whenever the `components.Footer` + // *function* identity changes -- a fresh arrow function here every render + // (streaming pushes a render per token) tore down and rebuilt the + // streaming bubble every frame. `components` itself must still get a new + // object each render (Virtuoso only redraws the slot when that reference + // changes), but Footer's own identity stays stable via the ref, so React + // reconciles the redraw as an update to the existing instance rather than + // an unmount/remount. + const trailingRef = useRef(trailingContent); + trailingRef.current = trailingContent; + const Footer = useCallback(() => <>{trailingRef.current}, []); + useImperativeHandle(ref, () => ({ scrollToBottom: () => { if (virtualized) { @@ -88,7 +100,7 @@ export const MessageList = forwardRef(function Message alignToBottom atBottomStateChange={onNearEndChange} itemContent={(index, message) => renderCard(message, index)} - components={{ Footer: () => <>{trailingContent} }} + components={{ Footer }} /> ); }); diff --git a/frontend/src/features/settings/SettingsPanel.test.tsx b/frontend/src/features/settings/SettingsPanel.test.tsx index 2fc1f59..6d577c4 100644 --- a/frontend/src/features/settings/SettingsPanel.test.tsx +++ b/frontend/src/features/settings/SettingsPanel.test.tsx @@ -66,6 +66,53 @@ describe("SettingsPanel", () => { })); }); + it("preserves the configured chat model when saving with an empty model inventory", async () => { + const user = userEvent.setup(); + const onSave = vi.fn<(settings: CortexSettings) => Promise>().mockResolvedValue(); + const settings: CortexSettings = { + appearance: { theme: "dark" }, + models: { chat: "gguf:mistral-7b.gguf", title: null, translation: "translategemma:4b" }, + generation: { temperature: 0.7, num_ctx: 4096, seed: -1, system_instructions: "" }, + }; + const models: ModelResponse = { + required_models: [], + optional_models: [], + installed_models: [], + models: [], + connection: { success: false, status: "error", message: "Ollama is not running." }, + }; + + render( + Promise>().mockResolvedValue()} + onReplaceMemory={vi.fn<(memos: string[]) => Promise>().mockResolvedValue()} + onClearMemory={vi.fn<() => Promise>().mockResolvedValue()} + models={models} + modelBusy={false} + modelProgress={null} + setupUrl="https://ollama.com/download" + onCheckModels={vi.fn<() => Promise>().mockResolvedValue()} + onPullModel={vi.fn<(model: string) => Promise>().mockResolvedValue()} + llamacppStatus={{ state: "idle", binary_present: false, loaded_model: null, last_error: null, models_directory: "" }} + onDownloadGGUF={vi.fn().mockResolvedValue(undefined)} + onClose={vi.fn()} + />, + ); + + // An unrelated edit, e.g. toggling the theme, must not wipe the still-valid + // configured chat model just because the inventory came back empty. + await user.click(screen.getByRole("button", { name: "Save settings" })); + + expect(onSave).toHaveBeenCalledWith(expect.objectContaining({ + models: expect.objectContaining({ chat: "gguf:mistral-7b.gguf", title: null }), + })); + }); + it("shows an active spinner and progress message while pulling the default translation model", async () => { const user = userEvent.setup(); const settings: CortexSettings = { diff --git a/frontend/src/features/settings/SettingsPanel.tsx b/frontend/src/features/settings/SettingsPanel.tsx index 083432e..2265228 100644 --- a/frontend/src/features/settings/SettingsPanel.tsx +++ b/frontend/src/features/settings/SettingsPanel.tsx @@ -101,7 +101,14 @@ export function SettingsPanel({ })); const saveDraft = () => onSave({ ...draft, - models: { ...modelSettings, chat: selectedChatModel || null, title: null }, + models: { + ...modelSettings, + // An empty inventory (Ollama down, a failed refresh) is a routine, + // recoverable state -- it must not overwrite a still-valid configured + // model with null just because the picker has nothing to offer right now. + chat: installedModels.length ? (selectedChatModel || null) : (modelSettings.chat ?? null), + title: null, + }, }); return ( diff --git a/main.py b/main.py index 56230ee..9ceda6a 100644 --- a/main.py +++ b/main.py @@ -177,7 +177,7 @@ def _monitor_native_window( server, readiness_url: str, ) -> None: - """Close the shell only after sustained backend-readiness failure.""" + """Close the shell only after sustained backend-liveness failure.""" failed_probes = 0 while not window.events.closed.is_set(): ready = wait_for_http( @@ -195,7 +195,7 @@ def _monitor_native_window( except Exception: pass raise RuntimeError("Cortex backend stopped unexpectedly.") from backend.error - if failed_probes >= 12: + if failed_probes >= 8: try: window.destroy() except Exception: @@ -203,7 +203,7 @@ def _monitor_native_window( if server.should_exit: return raise RuntimeError( - "Cortex backend became unavailable after 12 consecutive readiness probes." + "Cortex backend became unavailable after 8 consecutive liveness probes." ) if frontend is not None and not frontend.running: try: @@ -213,7 +213,7 @@ def _monitor_native_window( raise RuntimeError( f"Vite stopped unexpectedly with exit code {frontend.returncode}." ) - time.sleep(0.1) + time.sleep(1.5) def _run_headless(*, backend, frontend, server) -> int: @@ -352,7 +352,7 @@ def _run_web(args: argparse.Namespace) -> int: frontend=frontend, server=server, readiness_url=( - f"http://127.0.0.1:{backend_port}/api/v1/health/ready" + f"http://127.0.0.1:{backend_port}/api/v1/health/live" ), ), ) diff --git a/tests/test_api_contract.py b/tests/test_api_contract.py index 90ca83d..63365ff 100644 --- a/tests/test_api_contract.py +++ b/tests/test_api_contract.py @@ -97,6 +97,16 @@ def test_api_factory_is_headless_and_session_exchange_is_one_time(): assert second.status_code == 401 +def test_session_exchange_rejects_non_ascii_bootstrap_token_cleanly(): + app, client = _client() + with client: + response = client.post( + "/api/v1/session/exchange", + json={"bootstrap_token": "café-token"}, + ) + assert response.status_code == 401 + + def test_security_rejects_non_loopback_host_and_origin(): app, client = _client() default_app = create_app() @@ -146,6 +156,62 @@ def test_expired_session_is_rejected_without_exposing_token_details(): raise AssertionError("expired session was accepted") +def test_authenticate_slides_the_session_expiry_forward(): + """Regression guard: a session's expiry used to be fixed at issuance, so + the desktop app hard-locked after exactly one hour of continuous use -- + the frontend has no way to reach a fresh bootstrap token once its only + credential is destroyed after the initial handoff. Every successful + authenticate() must extend expires_at, so a session that is actually + being used never expires mid-session. + """ + manager = SessionManager(bootstrap_token="bootstrap", ttl_seconds=3600, allowed_hosts=("testserver",)) + exchanged = manager.exchange("bootstrap") + digest = manager._digest(exchanged.token) + # 50 minutes into a 60-minute TTL -- still valid, but would expire in + # 10 more minutes without a renewal. + stale_issued_at = datetime.now(timezone.utc) - timedelta(minutes=50) + manager._sessions[digest] = replace( + exchanged.principal, + issued_at=stale_issued_at, + expires_at=stale_issued_at + timedelta(seconds=3600), + ) + old_expiry = manager._sessions[digest].expires_at + + principal = manager.authenticate(exchanged.token) + + assert principal.expires_at > old_expiry + assert manager._sessions[digest].expires_at == principal.expires_at + # And the renewed session is genuinely usable well past the original + # one-hour mark, not just nominally not-yet-expired. + assert principal.expires_at > datetime.now(timezone.utc) + timedelta(minutes=55) + + +def test_authenticate_caps_the_sliding_expiry_at_the_absolute_max_lifetime(): + """A session cannot renew itself forever -- continuous use still hits + an absolute lifetime cap rather than sliding indefinitely.""" + manager = SessionManager( + bootstrap_token="bootstrap", + ttl_seconds=3600, + max_lifetime_seconds=7200, + allowed_hosts=("testserver",), + ) + exchanged = manager.exchange("bootstrap") + digest = manager._digest(exchanged.token) + issued_at = datetime.now(timezone.utc) - timedelta(hours=1, minutes=55) # close to the 2h cap + manager._sessions[digest] = replace( + exchanged.principal, + issued_at=issued_at, + # Not yet expired, but well below the eventual ~5-minute-away cap -- + # a realistic pre-renewal state, unlike setting it past the cap. + expires_at=datetime.now(timezone.utc) + timedelta(minutes=1), + ) + + principal = manager.authenticate(exchanged.token) + + assert principal.expires_at <= issued_at + timedelta(hours=2) + assert principal.expires_at > datetime.now(timezone.utc) + timedelta(minutes=1) + + def test_generation_selects_a_live_local_model_and_translation_is_opt_in(): settings = CortexSettings() snapshot = _generation_snapshot( @@ -624,3 +690,79 @@ def after_runner(sink, _cancel_event): assert after_registry.status(after.job_id, owner="owner").status == "succeeded" asyncio.run(exercise()) + + +def test_job_registry_shutdown_is_bounded_for_a_worker_that_never_observes_cancellation(): + """Regression guard: shutdown() used to await every pending worker with + no bound at all, including one stuck inside a synchronous call that + never polls cancel_event (a model HTTP request with no read deadline, + for example). That hung app shutdown -- and the llama-server child + process behind it -- for as long as that call took, sometimes forever. + A worker that has not begun committing its result must now be + abandoned once the grace period elapses so shutdown always completes + in bounded time. + """ + async def exercise(): + registry = JobRegistry(poll_seconds=0.001, shutdown_grace_seconds=0.05) + started = Event() + never_released = Event() + + def stuck_runner(_sink, _cancel_event): + # Never checks _cancel_event -- simulates a blocking call (e.g. + # a socket read with no deadline) that ignores cancellation. + started.set() + never_released.wait(timeout=5) + return {"persisted": False} + + job = await registry.start( + kind="generation", + owner="owner", + thread_id="thread-stuck", + runner=stuck_runner, + ) + for _ in range(200): + if started.is_set(): + break + await asyncio.sleep(0.001) + else: + raise AssertionError("worker did not start") + + loop = asyncio.get_event_loop() + started_at = loop.time() + await asyncio.wait_for(registry.shutdown(), timeout=1.0) + elapsed = loop.time() - started_at + + assert elapsed < 0.5, f"shutdown() took {elapsed:.2f}s, expected it bounded near the 0.05s grace period" + assert registry.status(job.job_id, owner="owner").status == "cancelling" + never_released.set() + + asyncio.run(exercise()) + + +def test_lifespan_runtime_teardown_runs_even_if_job_shutdown_raises(): + """Regression guard: the lifespan finally block used to await job + shutdown unconditionally before tearing down the runtime, so an + exception there (or, before the bounded-shutdown fix, an indefinite + hang) would skip llamacpp_manager.stop() entirely and leave the + llama-server child process orphaned. Runtime teardown must run + regardless of whether job shutdown succeeds. + """ + class _RaisingJobs: + async def shutdown(self): + raise RuntimeError("boom") + + class _FakeLlamaManager: + def __init__(self): + self.stopped = False + + def stop(self): + self.stopped = True + + fake_manager = _FakeLlamaManager() + app = create_app(allowed_hosts=ALLOWED_HOSTS, llamacpp_manager=fake_manager) + app.state.jobs = _RaisingJobs() + + with TestClient(app): + pass + + assert fake_manager.stopped is True diff --git a/tests/test_chat_client_routing.py b/tests/test_chat_client_routing.py index 2acbeef..d25125d 100644 --- a/tests/test_chat_client_routing.py +++ b/tests/test_chat_client_routing.py @@ -114,6 +114,61 @@ def chat(self, *, model, messages, options): assert result["options"] == {"temperature": 0.5} +def test_ollama_chat_client_stops_consuming_the_stream_once_cancelled() -> None: + """A cancellation_event switches OllamaChatClient to a streamed call it + can abort between chunks, closing the generator (which owns the + underlying httpx streaming response in the real ollama package -- see + Client._request) rather than reading it to completion.""" + from threading import Event + + closed = {"value": False} + + def chunk_generator(): + try: + yield {"message": {"content": "Hel"}, "done": False} + yield {"message": {"content": "lo"}, "done": False} + yield {"message": {}, "done": True, "prompt_eval_count": 5, "eval_count": 3} + except GeneratorExit: + closed["value"] = True + raise + + class _StubStreamingOllama: + def chat(self, *, model, messages, options, stream=False): + assert stream is True + return chunk_generator() + + already_cancelled = Event() + already_cancelled.set() + client = OllamaChatClient(_StubStreamingOllama()) + + result = client.chat(model="m", messages=[], options={}, cancellation_event=already_cancelled) + + assert result["message"]["content"] == "" + assert closed["value"] is True + + +def test_ollama_chat_client_streams_the_full_response_when_not_cancelled() -> None: + from threading import Event + + def chunk_generator(): + yield {"message": {"content": "Hel"}, "done": False} + yield {"message": {"content": "lo"}, "done": False} + yield {"message": {}, "done": True, "prompt_eval_count": 5, "eval_count": 3} + + class _StubStreamingOllama: + def chat(self, *, model, messages, options, stream=False): + assert stream is True + return chunk_generator() + + client = OllamaChatClient(_StubStreamingOllama()) + + result = client.chat(model="m", messages=[], options={}, cancellation_event=Event()) + + assert result["message"]["content"] == "Hello" + assert result["prompt_eval_count"] == 5 + assert result["eval_count"] == 3 + + class _StaticProvider: def __init__(self, base_url: str) -> None: self._base_url = base_url @@ -242,6 +297,66 @@ def test_a_runtime_fault_is_not_reported_as_a_refused_message() -> None: assert "could not accept" in client_message +def test_llamacpp_chat_client_stops_consuming_the_stream_once_cancelled(tmp_path: Path) -> None: + """Regression guard: chat() used to make one blocking, non-cancellable + request (stream: false), so Stop could not interrupt an in-flight call + until the model finished on its own -- up to the client's 600s read + timeout. Passing cancellation_event switches to the streamed request and + checks the event between chunks; an already-set event must stop the + client from consuming (and returning) any of the response. + """ + from threading import Event + + model_path = tmp_path / "tiny.gguf" + model_path.write_bytes(b"fake") + state = FakeLlamaCppState(generation_response="a long response that streams as several chunks") + app = create_fake_llamacpp_app(state) + http_client = TestClient(app, base_url="http://fakellama") + provider = _StaticProvider("http://fakellama") + client = LlamaCppChatClient(provider, models_directory=lambda: tmp_path, http_client=http_client) + + already_cancelled = Event() + already_cancelled.set() + + response = client.chat( + model=f"gguf:{model_path.name}", + messages=[{"role": "user", "content": "hi"}], + options={}, + cancellation_event=already_cancelled, + ) + + assert response["message"]["content"] == "" + + +def test_llamacpp_chat_client_streams_the_full_response_when_not_cancelled(tmp_path: Path) -> None: + """The streamed (cancellation_event given) and blocking (not given) + paths must produce the same adapted result when nothing is cancelled -- + passing an event that never fires should behave exactly like today's + ordinary call.""" + from threading import Event + + model_path = tmp_path / "tiny.gguf" + model_path.write_bytes(b"fake") + state = FakeLlamaCppState(generation_response="Hello from llama.cpp", generation_thoughts="pondering") + app = create_fake_llamacpp_app(state) + http_client = TestClient(app, base_url="http://fakellama") + provider = _StaticProvider("http://fakellama") + client = LlamaCppChatClient(provider, models_directory=lambda: tmp_path, http_client=http_client) + + response = client.chat( + model=f"gguf:{model_path.name}", + messages=[{"role": "user", "content": "hi"}], + options={"num_ctx": 4096, "temperature": 0.7}, + cancellation_event=Event(), + ) + + assert response["message"]["content"] == "Hello from llama.cpp" + assert response["message"]["thinking"] == "pondering" + assert response["prompt_eval_duration"] == 120_000_000 + assert response["eval_duration"] == 480_000_000 + assert response["eval_count"] == 48 + + def test_adapt_falls_back_to_wall_clock_when_timings_absent() -> None: payload = {"choices": [{"message": {"content": "hi"}}], "usage": {"prompt_tokens": 5, "completion_tokens": 3}} adapted = _adapt_to_ollama_shape(payload, elapsed_seconds=1.5) diff --git a/tests/test_chat_correctness.py b/tests/test_chat_correctness.py index b8bc6bc..483949f 100644 --- a/tests/test_chat_correctness.py +++ b/tests/test_chat_correctness.py @@ -168,6 +168,34 @@ def test_default_context_window_survives_a_realistic_long_conversation(self): "overhead and conversations will appear to lose their memory.", ) + def test_oversized_newest_exchange_does_not_wipe_the_rest_of_history(self): + """Regression guard: fit_history_to_context used to stop walking the + moment the single newest exchange alone exceeded the budget, discarding + every older exchange too and returning "No history available." even + though ten small exchanges right before it would easily have fit. The + newest exchange being oversized should just be dropped on its own. + """ + messages = [] + for index in range(10): + messages.append({"role": "user", "content": f"Question number {index} about the project"}) + messages.append({"role": "assistant", "content": f"Short answer number {index}."}) + messages.append({"role": "user", "content": "Please write the full module"}) + messages.append({"role": "assistant", "content": "X" * 35_000}) + + history = SynthesisAgent.fit_history_to_context( + messages, + query="now explain what you just did", + permanent_memories=[], + memories_enabled=True, + user_system_instructions=None, + num_ctx=8192, + ) + + self.assertNotEqual(history, "No history available.") + self.assertEqual(history.count("User: "), 10) + self.assertIn("Question number 9", history) + self.assertNotIn("X" * 100, history) + def test_context_budget_trims_oversized_permanent_memory(self): memories = [f"memory-{index} " + ("detail " * 120) for index in range(20)] diff --git a/tests/test_code_execution.py b/tests/test_code_execution.py index 08466a9..cc55255 100644 --- a/tests/test_code_execution.py +++ b/tests/test_code_execution.py @@ -112,6 +112,69 @@ def test_network_broker_rejects_private_targets() -> None: ) +def test_network_validation_returns_the_address_it_vetted(monkeypatch) -> None: + """The vetted address must come back so the caller can dial it directly. + + Discarding it and handing the hostname to the HTTP stack is what allowed + DNS rebinding: the stack resolved a second time, and a nameserver + answering differently on that second lookup reached targets the check + had just rejected. + """ + def fake_getaddrinfo(host, port, *args, **kwargs): + del host, port, args, kwargs + return [(0, 0, 0, "", ("93.184.216.34", 80))] + + monkeypatch.setattr(code_execution.socket, "getaddrinfo", fake_getaddrinfo) + + url, pinned_ip = code_execution._validate_network_url("http://example.test/status") + + assert url == "http://example.test/status" + assert pinned_ip == "93.184.216.34" + + +def test_network_broker_pins_the_vetted_address_against_dns_rebinding(monkeypatch) -> None: + """A nameserver that answers public-then-private must not win. + + The first resolution passes validation; a second resolution (the one the + HTTP stack would otherwise perform when it opens the socket) returns + loopback. The connection must still be made to the first, vetted address + -- never to the rebound one. + """ + answers = [ + [(0, 0, 0, "", ("93.184.216.34", 80))], # vetted: public + [(0, 0, 0, "", ("127.0.0.1", 80))], # rebound: loopback + ] + + def rebinding_getaddrinfo(host, port, *args, **kwargs): + del host, port, args, kwargs + return answers.pop(0) if len(answers) > 1 else answers[0] + + monkeypatch.setattr(code_execution.socket, "getaddrinfo", rebinding_getaddrinfo) + + _, pinned_ip = code_execution._validate_network_url("http://rebind.test/status") + assert pinned_ip == "93.184.216.34" + + dialed: list[tuple[str, int]] = [] + + def fake_create_connection(address, timeout=None, source_address=None): + del timeout, source_address + dialed.append(address) + raise OSError("connection not actually made in this test") + + plain, _tls = code_execution._pinned_connection_classes(pinned_ip) + connection = plain("rebind.test", 80) + monkeypatch.setattr(connection, "_create_connection", fake_create_connection) + with pytest.raises(OSError): + connection.connect() + + assert dialed == [("93.184.216.34", 80)], ( + "the connection re-resolved instead of using the vetted address" + ) + # The hostname is still what travels in Host / SNI, so servers and + # certificate validation are unaffected by the pinning. + assert connection.host == "rebind.test" + + def test_code_execution_waits_for_one_time_approval_and_returns_structured_output(tmp_path) -> None: repository = ExecutionRepository(tmp_path / "execution.sqlite", tmp_path / "artifacts") coordinator = LocalExecutionCoordinator(repository, code_timeout_seconds=3.0) diff --git a/tests/test_launcher.py b/tests/test_launcher.py index 9c5caab..9edbed7 100644 --- a/tests/test_launcher.py +++ b/tests/test_launcher.py @@ -311,6 +311,7 @@ def stop(self): self.running = False calls: list[tuple[str, object]] = [] + probed_urls: list[str] = [] monkeypatch.setattr(launcher_main, "InstanceLock", FakeInstance) monkeypatch.setattr(launcher_main, "_requested_port", lambda _port: 43125) monkeypatch.setattr(launcher_main, "ensure_frontend", lambda *_args, **_kwargs: tmp_path) @@ -318,7 +319,12 @@ def stop(self): monkeypatch.setattr(launcher_main, "_server_for_app", lambda *_args, **_kwargs: server) monkeypatch.setattr(launcher_main, "_install_shutdown_signals", lambda _server: None) monkeypatch.setattr(launcher_main, "ServerSupervisor", FakeBackend) - monkeypatch.setattr(launcher_main, "wait_for_http", lambda *_args, **_kwargs: True) + + def fake_wait_for_http(url, *_args, **_kwargs): + probed_urls.append(url) + return True + + monkeypatch.setattr(launcher_main, "wait_for_http", fake_wait_for_http) monkeypatch.setattr( launcher_main, "ensure_webview2_runtime", @@ -327,20 +333,76 @@ def stop(self): monkeypatch.setattr( launcher_main, "run_desktop_window", - lambda config, monitor: calls.append(("window", config)), + lambda config, monitor: calls.append(("window", (config, monitor))), ) args = launcher_main.build_parser().parse_args(["--data-dir", str(tmp_path)]) assert launcher_main._run_web(args) == 0 assert [name for name, _value in calls] == ["runtime", "window"] - window_config = calls[1][1] + window_config, monitor = calls[1][1] assert isinstance(window_config, DesktopWindowConfig) assert window_config.url == "http://127.0.0.1:43125/#bootstrap=bootstrap-token" assert window_config.storage_path == tmp_path / "webview" assert server.should_exit is True assert backend_instances[0].running is False + # Startup gate used the heavier readiness probe. + assert probed_urls == ["http://127.0.0.1:43125/api/v1/health/ready"] + + # The ongoing native-window monitor should poll the cheap liveness route + # rather than the readiness route, since it runs for the app's lifetime. + closed_checks = {"count": 0} + + def closed_is_set() -> bool: + closed_checks["count"] += 1 + return closed_checks["count"] > 1 + + fake_window = SimpleNamespace( + events=SimpleNamespace(closed=SimpleNamespace(is_set=closed_is_set)), + destroy=lambda: None, + ) + monkeypatch.setattr(launcher_main.time, "sleep", lambda *_args, **_kwargs: None) + monitor(fake_window) + assert probed_urls[-1] == "http://127.0.0.1:43125/api/v1/health/live" + + +def test_monitor_native_window_polls_slowly_and_grants_a_multi_second_grace_period( + monkeypatch: pytest.MonkeyPatch, +): + sleeps: list[float] = [] + monkeypatch.setattr(launcher_main.time, "sleep", lambda seconds: sleeps.append(seconds)) + + probed_urls: list[str] = [] + + def fake_wait_for_http(url, *, timeout, is_alive): + probed_urls.append(url) + return False + + monkeypatch.setattr(launcher_main, "wait_for_http", fake_wait_for_http) + + window = SimpleNamespace( + events=SimpleNamespace(closed=SimpleNamespace(is_set=lambda: False)), + destroy=lambda: destroyed.append(True), + ) + destroyed: list[bool] = [] + backend = SimpleNamespace(error=None) + frontend = SimpleNamespace(running=True) + server = SimpleNamespace(should_exit=False) + + with pytest.raises(RuntimeError, match="8 consecutive liveness probes"): + launcher_main._monitor_native_window( + window, + backend=backend, + frontend=frontend, + server=server, + readiness_url="http://127.0.0.1:43125/api/v1/health/live", + ) + + assert probed_urls == ["http://127.0.0.1:43125/api/v1/health/live"] * 8 + assert sleeps == [1.5] * 7 + assert destroyed == [True] + def _frontend_fixture(tmp_path: Path) -> Path: root = tmp_path / "frontend" @@ -455,7 +517,9 @@ def test_frontend_build_stages_sources_outside_live_node_modules( installed_roots: list[Path] = [] monkeypatch.setattr(frontend_module, "_major_version", lambda _: 24) - def fake_install(frontend_root: Path, _expected_lock_digest: str) -> None: + def fake_install( + frontend_root: Path, _expected_lock_digest: str, _cache_root: Path + ) -> None: installed_roots.append(frontend_root) def fake_run(command: list[str], *, cwd: Path) -> None: @@ -508,6 +572,129 @@ def fake_run(command: list[str], *, cwd: Path) -> None: assert not list(tmp_path.glob(".cortex-frontend-build-*")) +def test_stale_staging_directories_are_swept_before_a_new_build( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + root = _frontend_fixture(tmp_path) + stale = tmp_path / ".cortex-frontend-build-orphaned" + stale.mkdir() + (stale / "leftover.txt").write_text("orphaned", encoding="utf-8") + monkeypatch.setattr(frontend_module, "_major_version", lambda _: 24) + monkeypatch.setattr(frontend_module, "_install_if_needed", lambda *_args: None) + + def fake_run(command: list[str], *, cwd: Path) -> None: + staging = Path(command[-1]) + staging.mkdir(parents=True) + (staging / "index.html").write_text("new", encoding="utf-8") + + monkeypatch.setattr(frontend_module, "_run", fake_run) + + frontend_module.build_frontend(root) + + assert not stale.exists() + assert not list(tmp_path.glob(".cortex-frontend-build-*")) + + +def test_reclaim_stale_staging_directories_removes_orphaned_builds(tmp_path: Path): + stale_a = tmp_path / ".cortex-frontend-build-aaa" + stale_b = tmp_path / ".cortex-frontend-build-bbb" + keep = tmp_path / ".cortex-frontend-build-new" + stale_a.mkdir() + (stale_a / "leftover.txt").write_text("orphaned", encoding="utf-8") + stale_b.mkdir() + + frontend_module._reclaim_stale_staging_directories(tmp_path, keep) + + assert not stale_a.exists() + assert not stale_b.exists() + + +def test_stale_staging_directory_removal_failure_is_logged_not_raised( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +): + locked = tmp_path / ".cortex-frontend-build-locked" + locked.mkdir() + + def flaky_rmtree(_path, *_args, **_kwargs): + raise OSError("file is locked by another process") + + monkeypatch.setattr(frontend_module.shutil, "rmtree", flaky_rmtree) + + with caplog.at_level("WARNING"): + frontend_module._reclaim_stale_staging_directories( + tmp_path, tmp_path / ".cortex-frontend-build-new" + ) + + assert locked.exists() + assert "Could not remove stale frontend build directory" in caplog.text + + +def test_install_cache_hit_skips_npm_ci_for_unchanged_lockfile( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + build_root = tmp_path / "build" + build_root.mkdir() + cache_root = tmp_path / "cache" + cached_modules = cache_root / "node_modules" + cached_modules.mkdir(parents=True) + (cached_modules / "package.json").write_text("{}", encoding="utf-8") + (cache_root / frontend_module.INSTALL_MANIFEST_NAME).write_text( + json.dumps({"lock_digest": "abc123"}), encoding="utf-8" + ) + + def fail_run(*_args, **_kwargs): + pytest.fail("npm ci should not run on a cache hit") + + monkeypatch.setattr(frontend_module, "_run", fail_run) + + frontend_module._install_if_needed(build_root, "abc123", cache_root) + + assert (build_root / "node_modules" / "package.json").read_text(encoding="utf-8") == "{}" + + +def test_install_cache_miss_runs_npm_ci_when_lockfile_changes( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + build_root = tmp_path / "build" + build_root.mkdir() + cache_root = tmp_path / "cache" + cached_modules = cache_root / "node_modules" + cached_modules.mkdir(parents=True) + (cached_modules / "package.json").write_text('{"old": true}', encoding="utf-8") + (cache_root / frontend_module.INSTALL_MANIFEST_NAME).write_text( + json.dumps({"lock_digest": "old-digest"}), encoding="utf-8" + ) + + calls: list[Path] = [] + + def fake_run(command: list[str], *, cwd: Path) -> None: + calls.append(cwd) + node_modules = cwd / "node_modules" + node_modules.mkdir(parents=True) + (node_modules / "package.json").write_text('{"new": true}', encoding="utf-8") + + monkeypatch.setattr(frontend_module, "_run", fake_run) + + frontend_module._install_if_needed(build_root, "new-digest", cache_root) + + assert calls == [build_root] + marker_path = cache_root / frontend_module.INSTALL_MANIFEST_NAME + stored = json.loads(marker_path.read_text(encoding="utf-8")) + assert stored["lock_digest"] == "new-digest" + assert (cached_modules / "package.json").read_text(encoding="utf-8") == '{"new": true}' + + # A later build with the same lockfile digest hits the refreshed cache. + calls.clear() + build_root_2 = tmp_path / "build2" + build_root_2.mkdir() + frontend_module._install_if_needed(build_root_2, "new-digest", cache_root) + + assert calls == [] + assert ( + build_root_2 / "node_modules" / "package.json" + ).read_text(encoding="utf-8") == '{"new": true}' + + def test_frontend_install_failure_leaves_live_node_modules_untouched( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ): @@ -600,3 +787,17 @@ def test_handoff_rotates_bootstrap_token_and_shutdown_is_authenticated(): assert shutdown.json() == {"status": "accepted"} assert shutdown_calls == [True] assert client.get("/api/v1/health/ready").status_code == 503 + + +def test_handoff_rejects_non_ascii_header_with_a_clean_unauthorized(): + app = create_app( + build_demo_dependencies(), + allowed_hosts=("testserver", "127.0.0.1", "localhost", "::1"), + handoff_secret="handoff-secret", + ) + with TestClient(app) as client: + response = client.post( + "/api/v1/session/handoff", + headers={b"X-Cortex-Handoff": "café-token".encode("latin-1")}, + ) + assert response.status_code == 401 diff --git a/tests/test_llamacpp_binary_fetcher.py b/tests/test_llamacpp_binary_fetcher.py index 4e87d0b..9ee7264 100644 --- a/tests/test_llamacpp_binary_fetcher.py +++ b/tests/test_llamacpp_binary_fetcher.py @@ -8,7 +8,9 @@ import hashlib import io +import os import tempfile +import time import zipfile from pathlib import Path from unittest.mock import patch @@ -158,12 +160,44 @@ def test_is_cached_reports_false_instead_of_raising_when_hashing_hits_memory_pre archive_bytes = _build_archive() release = _release_for(archive_bytes) fetcher = BinaryFetcher(tmp_path, http_client=_client_returning(archive_bytes)) - fetcher.ensure_binary(release, "cpu") + exe_path = fetcher.ensure_binary(release, "cpu") + + # Touch a file so the cheap tree-identity cache misses and this call + # actually reaches hash_directory -- otherwise the unchanged-directory + # fast path would return the cached result without calling it at all. + future = time.time() + 10 + os.utime(exe_path, (future, future)) with patch("cortex_backend.llamacpp.binary_fetcher.hash_directory", side_effect=MemoryError): assert fetcher.is_cached(release, "cpu") is False +def test_is_cached_only_hashes_once_for_an_unchanged_directory(tmp_path: Path) -> None: + """/api/v1/system polls is_cached() every ~2s while idle -- the full + SHA-256 directory walk must be skipped when nothing on disk changed, and + only re-run once a file actually changes.""" + archive_bytes = _build_archive() + release = _release_for(archive_bytes) + fetcher = BinaryFetcher(tmp_path, http_client=_client_returning(archive_bytes)) + exe_path = fetcher.ensure_binary(release, "cpu") + + with patch( + "cortex_backend.llamacpp.binary_fetcher.hash_directory", wraps=hash_directory + ) as spy: + # ensure_binary() already verified (and cached) this directory, so + # a repeated is_cached() call for the same unchanged tree must not + # re-hash it. + assert fetcher.is_cached(release, "cpu") is True + assert fetcher.is_cached(release, "cpu") is True + assert spy.call_count == 0 + + future = time.time() + 10 + os.utime(exe_path, (future, future)) + + assert fetcher.is_cached(release, "cpu") is True + assert spy.call_count == 1 + + def test_zip_slip_entries_are_rejected(tmp_path: Path) -> None: buffer = io.BytesIO() with zipfile.ZipFile(buffer, "w") as archive: diff --git a/tests/test_llamacpp_server_manager.py b/tests/test_llamacpp_server_manager.py index 967e956..6c58181 100644 --- a/tests/test_llamacpp_server_manager.py +++ b/tests/test_llamacpp_server_manager.py @@ -248,7 +248,9 @@ def test_vulkan_failure_falls_back_to_cpu_and_is_cached(tmp_path: Path) -> None: assert handle is not None assert fetcher.ensure_binary_calls == ["vulkan", "cpu"] marker = json.loads((tmp_path / "preferred_gpu_backend.json").read_text("utf-8")) - assert marker == {"known_bad": "vulkan"} + assert marker["known_bad"] == "vulkan" + assert marker["model"] == str(tmp_path / "model.gguf") + assert marker["num_ctx"] == 4096 # A fresh manager instance (simulating an app restart) must read the # cached marker and go straight to cpu -- no repeated failed attempt. @@ -261,6 +263,67 @@ def test_vulkan_failure_falls_back_to_cpu_and_is_cached(tmp_path: Path) -> None: assert fetcher2.ensure_binary_calls == ["cpu"] +def test_known_bad_backend_marker_does_not_affect_a_different_model_or_context(tmp_path: Path) -> None: + """Regression guard: the known-bad marker used to be a single global + string, so one oversized model failing on vulkan permanently pushed + every other model -- and every other context size -- to cpu too. The + marker must be scoped to the exact (model, num_ctx) that failed. + """ + fetcher = _FakeFetcher() + launcher = _QueueLauncher([_FakePopen(exit_immediately=True), _FakePopen()]) + manager = _manager( + tmp_path, fetcher=fetcher, launcher=launcher, http_client=_AlwaysHealthyClient(), gpu_backend="auto" + ) + manager.ensure_ready(tmp_path / "big-model.gguf", num_ctx=8192) + assert fetcher.ensure_binary_calls == ["vulkan", "cpu"] + + # A different model must still be tried on vulkan first. + fetcher2 = _FakeFetcher() + launcher2 = _QueueLauncher([_FakePopen()]) + manager2 = _manager( + tmp_path, fetcher=fetcher2, launcher=launcher2, http_client=_AlwaysHealthyClient(), gpu_backend="auto" + ) + manager2.ensure_ready(tmp_path / "small-model.gguf", num_ctx=8192) + assert fetcher2.ensure_binary_calls == ["vulkan"] + + # The same model at a different context size must also still be tried + # on vulkan first -- a smaller context is exactly the kind of change + # that can make an otherwise-too-large model fit. + fetcher3 = _FakeFetcher() + launcher3 = _QueueLauncher([_FakePopen()]) + manager3 = _manager( + tmp_path, fetcher=fetcher3, launcher=launcher3, http_client=_AlwaysHealthyClient(), gpu_backend="auto" + ) + manager3.ensure_ready(tmp_path / "big-model.gguf", num_ctx=2048) + assert fetcher3.ensure_binary_calls == ["vulkan"] + + +def test_known_bad_backend_marker_expires(tmp_path: Path) -> None: + """A stale marker (past the TTL) must not permanently pin cpu -- a + driver update or freed VRAM deserves a retry rather than an indefinite, + unrecoverable-without-manual-intervention ban.""" + marker_path = tmp_path / "preferred_gpu_backend.json" + marker_path.write_text( + json.dumps({ + "known_bad": "vulkan", + "model": str(tmp_path / "model.gguf"), + "num_ctx": 4096, + "release": None, + "at": time.time() - (25 * 3600), # older than the 24h TTL + }), + encoding="utf-8", + ) + fetcher = _FakeFetcher() + launcher = _QueueLauncher([_FakePopen()]) + manager = _manager( + tmp_path, fetcher=fetcher, launcher=launcher, http_client=_AlwaysHealthyClient(), gpu_backend="auto" + ) + + manager.ensure_ready(tmp_path / "model.gguf", num_ctx=4096) + + assert fetcher.ensure_binary_calls == ["vulkan"] + + def test_explicit_backend_setting_skips_fallback(tmp_path: Path) -> None: fetcher = _FakeFetcher() launcher = _QueueLauncher([_FakePopen(exit_immediately=True)]) @@ -487,6 +550,79 @@ def test_a_crash_loop_stops_with_an_honest_error_instead_of_thrashing(tmp_path: assert manager.status.state == "ready" +class _SlowTerminatePopen: + """Takes real wall-clock time to exit after terminate(), so a test can + observe whether something else was blocked meanwhile.""" + + def __init__(self, *, delay_seconds: float) -> None: + self._delay_seconds = delay_seconds + self.terminated = False + + def poll(self): + return None + + def terminate(self) -> None: + self.terminated = True + + def kill(self) -> None: + pass + + def wait(self, timeout=None): + time.sleep(self._delay_seconds) + return 0 + + +def test_crash_loop_guard_termination_does_not_block_status_polls(tmp_path: Path) -> None: + """Regression guard: the crash-loop guard used to tear the process down + while still holding the state lock the class documents as held for + microseconds only, so the runtime-status endpoint (polled every couple + of seconds by the UI) froze for the whole grace wait right as the guard + fired to report the honest "does not fit in memory" error. + """ + fetcher = _FakeFetcher() + manager = _manager(tmp_path, fetcher=fetcher, launcher=_QueueLauncher([]), http_client=_AlwaysHealthyClient()) + model_path = tmp_path / "model.gguf" + + # Arm the guard directly with a process that takes real time to exit, + # rather than driving three full crash/relaunch cycles just to get one + # in place -- what's under test is the guard's own teardown, not the + # counting that leads up to it (covered above). + slow_process = _SlowTerminatePopen(delay_seconds=0.3) + with manager._state_lock: + manager._process = slow_process + manager._loaded_model_path = model_path + manager._loaded_num_ctx = 6144 + manager._failure_key = (model_path, 6144) + manager._failure_times = [time.monotonic()] * 3 + manager._last_restart_reason = "simulated crash" + + max_poll_latency = 0.0 + stop_polling = threading.Event() + + def poll_status() -> None: + nonlocal max_poll_latency + while not stop_polling.is_set(): + started = time.monotonic() + _ = manager.status.state + max_poll_latency = max(max_poll_latency, time.monotonic() - started) + time.sleep(0.01) + + poller = threading.Thread(target=poll_status, daemon=True) + poller.start() + time.sleep(0.03) # let the poller get going before the guard fires + + with pytest.raises(LlamaCppError): + manager._guard_against_crash_loop(model_path, 6144) + + stop_polling.set() + poller.join(timeout=2.0) + + assert slow_process.terminated + assert max_poll_latency < 0.15, ( + f"a status poll took {max_poll_latency:.3f}s -- the state lock was held during termination" + ) + + def test_status_stays_responsive_while_a_model_loads(tmp_path: Path) -> None: """The UI polls status every couple of seconds. It must never queue behind a model load, which can legitimately take minutes.""" @@ -525,3 +661,134 @@ def __call__(self, argv: list[str], *, cwd: Path): assert not worker.is_alive() assert observed_starting is True assert manager.status.state == "ready" + + +class _FakeProcessWithPid: + def __init__(self, pid: int) -> None: + self.pid = pid + + +class _FakeWin32Job: + """Records the kernel32 call sequence without touching real Windows APIs.""" + + def __init__(self, *, create_job_result: int = 1, set_info_result: bool = True) -> None: + self.create_job_result = create_job_result + self.set_info_result = set_info_result + self.calls: list[tuple] = [] + self._next_handle = 100 + + def CreateJobObjectW(self, security_attributes, name): + self.calls.append(("CreateJobObjectW",)) + if not self.create_job_result: + return 0 + self._next_handle += 1 + return self._next_handle + + def SetInformationJobObject(self, job, info_class, info, info_size): + from cortex_backend.llamacpp.server_manager import _JobObjectExtendedLimitInformation + import ctypes as _ctypes + + limits = _ctypes.cast(info, _ctypes.POINTER(_JobObjectExtendedLimitInformation)).contents + self.calls.append(( + "SetInformationJobObject", + job, + info_class, + limits.basic_limit_information.limit_flags, + info_size, + )) + return 1 if self.set_info_result else 0 + + def OpenProcess(self, access, inherit_handle, pid): + self._next_handle += 1 + handle = self._next_handle + self.calls.append(("OpenProcess", access, inherit_handle, pid, handle)) + return handle + + def AssignProcessToJobObject(self, job, process): + self.calls.append(("AssignProcessToJobObject", job, process)) + return 1 + + def CloseHandle(self, handle): + self.calls.append(("CloseHandle", handle)) + return 1 + + +def test_job_object_launcher_applies_kill_on_close_policy_and_reuses_the_job(): + """Regression guard: llama-server was launched with no Job Object at + all, so any hard exit of Cortex (Task Manager, a crash) left it running + and holding the model resident. The launcher must create a Job Object + with JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE, assign each launched process to + it, and reuse the same job across restarts rather than leaking a handle + per relaunch. + """ + from cortex_backend.llamacpp.server_manager import ( + _JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE, + _PROCESS_SET_QUOTA, + _PROCESS_TERMINATE, + _JobObjectLauncher, + ) + + fake_win32 = _FakeWin32Job() + launcher = _JobObjectLauncher(win32_factory=lambda: fake_win32) + + launcher._apply_job_policy(_FakeProcessWithPid(pid=4242)) + + create_calls = [call for call in fake_win32.calls if call[0] == "CreateJobObjectW"] + assert len(create_calls) == 1 + set_info_calls = [call for call in fake_win32.calls if call[0] == "SetInformationJobObject"] + assert len(set_info_calls) == 1 + assert set_info_calls[0][3] == _JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE + open_calls = [call for call in fake_win32.calls if call[0] == "OpenProcess"] + assert open_calls == [("OpenProcess", _PROCESS_SET_QUOTA | _PROCESS_TERMINATE, False, 4242, open_calls[0][4])] + assign_calls = [call for call in fake_win32.calls if call[0] == "AssignProcessToJobObject"] + assert len(assign_calls) == 1 + job_handle = launcher._job + assert assign_calls[0] == ("AssignProcessToJobObject", job_handle, open_calls[0][4]) + # The process handle opened just to assign the job is closed again. + close_calls = [call for call in fake_win32.calls if call[0] == "CloseHandle"] + assert close_calls == [("CloseHandle", open_calls[0][4])] + + launcher._apply_job_policy(_FakeProcessWithPid(pid=5555)) + + create_calls = [call for call in fake_win32.calls if call[0] == "CreateJobObjectW"] + assert len(create_calls) == 1, "the job must be reused, not recreated, on a second launch" + assign_calls = [call for call in fake_win32.calls if call[0] == "AssignProcessToJobObject"] + assert len(assign_calls) == 2 + assert assign_calls[1][1] == job_handle + + +def test_job_object_launcher_does_not_break_startup_if_job_creation_fails(): + """A Job Object is defense in depth, not a hard requirement -- if kernel32 + refuses (sandboxed environment, exhausted handle quota, anything), the + local model runtime must still start normally rather than the failure + propagating and blocking generation entirely.""" + from cortex_backend.llamacpp.server_manager import _JobObjectLauncher + + fake_win32 = _FakeWin32Job(create_job_result=0) + launcher = _JobObjectLauncher(win32_factory=lambda: fake_win32) + + launcher._apply_job_policy(_FakeProcessWithPid(pid=1)) # must not raise + + assert launcher._job is None + assert not any(call[0] == "SetInformationJobObject" for call in fake_win32.calls) + + +def test_job_object_launcher_discards_a_misconfigured_job_instead_of_reusing_it(): + from cortex_backend.llamacpp.server_manager import _JobObjectLauncher + + fake_win32 = _FakeWin32Job(set_info_result=False) + launcher = _JobObjectLauncher(win32_factory=lambda: fake_win32) + + launcher._apply_job_policy(_FakeProcessWithPid(pid=1)) # must not raise + + assert launcher._job is None + close_calls = [call for call in fake_win32.calls if call[0] == "CloseHandle"] + assert len(close_calls) == 1, "the unusable job handle must be closed, not leaked" + assert not any(call[0] == "AssignProcessToJobObject" for call in fake_win32.calls) + + +def test_default_launcher_is_a_job_object_launcher(): + from cortex_backend.llamacpp.server_manager import _JobObjectLauncher, default_launcher + + assert isinstance(default_launcher, _JobObjectLauncher) + assert callable(default_launcher) diff --git a/tests/test_persistence.py b/tests/test_persistence.py index a2feb35..d786c76 100644 --- a/tests/test_persistence.py +++ b/tests/test_persistence.py @@ -112,6 +112,31 @@ def test_migration_migrates_skips_and_quarantines_per_file(self): self.assertTrue((legacy / "quarantine" / "malformed.json").exists()) self.assertTrue(list(root.glob("legacy_migrated_*/*.json"))) + def test_permanent_memory_add_memo_is_safe_across_threads(self): + with tempfile.TemporaryDirectory() as directory: + memory_path = Path(directory) / "memory_bank.json" + manager = PermanentMemoryManager(memory_file_path=str(memory_path)) + errors = [] + + def add_memo(index): + try: + manager.add_memo(f"memo {index}") + except Exception as exc: # pragma: no cover - assertion below reports it + errors.append(exc) + + threads = [threading.Thread(target=add_memo, args=(index,)) for index in range(20)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + + self.assertEqual(errors, []) + expected = {f"memo {index}" for index in range(20)} + self.assertEqual(set(manager.get_memos()), expected) + + reloaded = PermanentMemoryManager(memory_file_path=str(memory_path)) + self.assertEqual(set(reloaded.get_memos()), expected) + def test_permanent_memory_recovers_from_backup_after_interrupted_write(self): with tempfile.TemporaryDirectory() as directory: memory_path = Path(directory) / "memory_bank.json" diff --git a/tests/test_services.py b/tests/test_services.py index b8163b0..d3ca59e 100644 --- a/tests/test_services.py +++ b/tests/test_services.py @@ -10,6 +10,7 @@ from cortex_backend.core.generation import ( CodeExecutionProposal, + GenerationAttachment, GenerationSnapshot, GenerationStats, MemoryCommand, @@ -17,6 +18,7 @@ TranslationResult, ) from cortex_backend.services.generation import GenerationService +from cortex_backend.services.llm import SynthesisAgent from cortex_backend.services.models import ModelService from cortex_backend.services.progress import ProgressEvent @@ -165,6 +167,71 @@ def test_generation_is_headless_and_emits_owned_typed_progress(self): self.assertEqual(engine.memory_inputs, ["remember tea"]) self.assertEqual(engine.options["num_ctx"], 4096) + def test_attachment_reference_text_is_not_crushed_by_a_full_history_fit(self): + """Regression guard: history used to be fit to the context budget + first, greedily claiming nearly all of it before attachments were + ever considered, so a document attached mid-conversation could be + cut to a tiny fragment even though it would easily have fit had it + been given any priority over old chat turns. Attachments must now be + reserved room before history is sized around them. + """ + class _CapturingChatClient: + def __init__(self): + self.last_messages: list[dict] | None = None + + def chat(self, *, model, messages, options): + del model, options + self.last_messages = messages + return {"message": {"content": "ok", "thinking": None}} + + history_messages = [] + for index in range(20): + history_messages.append({"role": "user", "content": f"old-{index} " + ("details " * 80)}) + history_messages.append({"role": "assistant", "content": f"reply-{index} " + ("context " * 80)}) + + attachment = GenerationAttachment( + attachment_id="doc-1", + filename="report.md", + mime_type="text/markdown", + kind="document", + text_content="report line " * 400, + ) + + client = _CapturingChatClient() + service = GenerationService( + history_loader=lambda thread_id: history_messages, + memory_loader=lambda: [], + engine_factory=lambda snapshot: SynthesisAgent( + "chat-model", "title-model", "translate-model", client, + ), + ) + snapshot = GenerationSnapshot( + job_id="job-1", + thread_id="thread-1", + user_input="Summarize the attached report.", + model="chat-model", + title_model="title-model", + translation_model="translate-model", + model_options={"temperature": 0.7, "num_ctx": 4096, "seed": -1}, + memories_enabled=False, + translation_enabled=False, + target_language="French", + user_system_instructions=None, + attachments=(attachment,), + ) + + service.generate(snapshot) + + self.assertIsNotNone(client.last_messages) + sent_content = "\n".join(str(message.get("content", "")) for message in client.last_messages or []) + retained_repeats = sent_content.count("report line") + self.assertGreater( + retained_repeats, + 300, + f"Only {retained_repeats}/400 repetitions of the attachment text survived context " + "fitting -- the attachment was crushed by history claiming the whole budget first.", + ) + def test_engine_status_callback_reports_as_loading_model_progress(self): """An engine backed by a locally-managed runtime (llama.cpp) can report its own startup progress through the normal progress sink, diff --git a/tests/test_stage5_system.py b/tests/test_stage5_system.py index f62b073..4754447 100644 --- a/tests/test_stage5_system.py +++ b/tests/test_stage5_system.py @@ -8,8 +8,11 @@ import shutil from threading import Barrier +import httpx +import pytest from fastapi.testclient import TestClient +import Cortex_Preview from Cortex_Preview import build_preview_app from cortex_backend.api import build_demo_dependencies, create_app from cortex_backend.repositories.legacy_settings import LegacySettingsReader @@ -205,6 +208,25 @@ def test_packaged_runtime_builder_opens_existing_chat_fixture_without_qt(tmp_pat assert list(tmp_path.glob("chat_history_migrated_*/*.json")) +def test_preview_app_builds_ollama_client_with_a_bounded_timeout( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +): + captured_kwargs: dict[str, object] = {} + real_client_cls = Cortex_Preview.ollama.Client + + class RecordingClient(real_client_cls): + def __init__(self, *, host=None, **kwargs): + captured_kwargs.update(kwargs) + super().__init__(host=host, **kwargs) + + monkeypatch.setattr(Cortex_Preview.ollama, "Client", RecordingClient) + + build_preview_app(data_dir=tmp_path, serve_frontend=False) + + assert captured_kwargs.get("timeout") is not None + assert isinstance(captured_kwargs["timeout"], httpx.Timeout) + + def test_model_inventory_pull_progress_and_failure_are_safe(): state = FakeOllamaState(installed_models={"qwen3:8b"}) app = create_app( @@ -301,3 +323,77 @@ def test_diagnostics_exposes_migration_and_setup_capabilities(): payload = diagnostics.json() assert payload["settings_source"] == "memory" assert payload["ollama_setup_url"] == "https://ollama.com/download" + + +def test_settings_live_in_their_own_database_not_the_chat_database(tmp_path: Path): + """Settings writes take a full-file backup copy first, so colocating them + with chat history meant every settings save byte-copied the whole + transcript store. The settings database must be a separate file, and its + backup must never touch the chat database.""" + chat_db = tmp_path / "cortex_db.sqlite" + settings_db = tmp_path / "cortex_settings.sqlite" + chat_db.write_bytes(b"pretend this is a large chat history" * 1000) + chat_before = chat_db.read_bytes() + + repository = SQLiteSettingsRepository(settings_db, adopt_from=chat_db) + saved = repository.load().settings + repository.save(saved) + + assert settings_db.exists() + # The backup that save() takes is of the settings file, not the chat one. + assert repository.backup_path == Path(f"{settings_db}.bak") + assert not Path(f"{chat_db}.bak").exists() + assert chat_db.read_bytes() == chat_before + + +def test_settings_colocated_in_the_chat_database_are_adopted_once(tmp_path: Path): + """Regression guard for the upgrade path: without adoption, every existing + install would silently revert to default settings the first time it ran a + build that moved settings into their own file.""" + chat_db = tmp_path / "cortex_db.sqlite" + settings_db = tmp_path / "cortex_settings.sqlite" + + # An install from before the split: settings living inside the chat database. + old = SQLiteSettingsRepository(chat_db) + configured = old.load().settings + configured = configured.model_copy( + update={"models": configured.models.model_copy(update={"chat": "gguf:kept.gguf"})} + ) + old.save(configured) + + adopted = SQLiteSettingsRepository(settings_db, adopt_from=chat_db) + + assert adopted.load().settings.models.chat == "gguf:kept.gguf" + with adopted.connect() as connection: + assert connection.execute( + "SELECT COUNT(*) FROM settings_migration_ledger WHERE migration_key = ?", + ("chatdb-colocated-settings-to-own-file-v1",), + ).fetchone()[0] == 1 + + # Adoption is one-time: a later edit is not clobbered by re-adopting the + # stale row on the next startup. + current = adopted.load().settings + adopted.save( + current.model_copy( + update={"models": current.models.model_copy(update={"chat": "gguf:newer.gguf"})} + ) + ) + reopened = SQLiteSettingsRepository(settings_db, adopt_from=chat_db) + assert reopened.load().settings.models.chat == "gguf:newer.gguf" + + +def test_adoption_is_skipped_cleanly_when_there_is_nothing_to_adopt(tmp_path: Path): + """A fresh install, and an unreadable or settings-free old database, must + both start normally rather than failing closed.""" + settings_db = tmp_path / "cortex_settings.sqlite" + + # Nothing at the old path at all. + fresh = SQLiteSettingsRepository(settings_db, adopt_from=tmp_path / "absent.sqlite") + assert fresh.load().settings is not None + + # An old database that exists but holds no settings table. + other_db = tmp_path / "unrelated.sqlite" + other_db.write_bytes(b"not a sqlite file at all") + second = tmp_path / "second_settings.sqlite" + repository = SQLiteSettingsRepository(second, adopt_from=other_db) + assert repository.load().settings is not None