From 40160b68df6ecc464e398adc03a56c642eaa5ed2 Mon Sep 17 00:00:00 2001 From: AronAxe <33731256+AronAxe@users.noreply.github.com> Date: Tue, 25 Aug 2026 08:35:35 +0200 Subject: [PATCH] feat: add async runtime and concurrent vault support --- CHANGELOG.md | 6 + README.md | 37 +++- src/rtk_hermes_plus/__init__.py | 10 +- src/rtk_hermes_plus/async_runtime.py | 242 +++++++++++++++++++++ src/rtk_hermes_plus/cancellation.py | 63 ++++++ src/rtk_hermes_plus/compress.py | 114 +++++++++- src/rtk_hermes_plus/plugin.py | 46 +++- src/rtk_hermes_plus/rewrite.py | 121 ++++++++++- src/rtk_hermes_plus/storage.py | 15 +- tests/test_async_runtime.py | 314 +++++++++++++++++++++++++++ tests/test_storage_concurrency.py | 113 ++++++++++ 11 files changed, 1052 insertions(+), 29 deletions(-) create mode 100644 src/rtk_hermes_plus/async_runtime.py create mode 100644 src/rtk_hermes_plus/cancellation.py create mode 100644 tests/test_async_runtime.py create mode 100644 tests/test_storage_concurrency.py diff --git a/CHANGELOG.md b/CHANGELOG.md index d8c02c0..baf33ad 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,11 @@ # Changelog +## Unreleased + +- Added an `AsyncRuntime` façade and reusable thread-safe `CancellationToken` without changing the synchronous runtime API. +- Added cancellable asyncio RTK command rewriting and aggressive reads that kill and reap subprocesses on task or token cancellation. +- Enabled WAL-backed concurrent artifact-vault access and added deterministic thread/process contention and lease coverage. + ## 0.3.1 - 2026-08-24 - Kept collapsed-turn summaries inside the first retained user message so strict provider role sequencing remains valid. diff --git a/README.md b/README.md index aa307a1..fb0c485 100644 --- a/README.md +++ b/README.md @@ -29,6 +29,8 @@ The engine has four cooperating reduction paths: The reduction core is not intrinsically tied to Hermes: it operates on Python dictionaries, strings, stable request/session identifiers, and a local SQLite vault. The repository includes a turnkey Hermes plugin because Hermes exposes the required lifecycle hooks. Other agent runtimes need a small adapter that presents the same boundaries; they do not need a fork of the reduction engine. +Async agent frameworks can use the included `AsyncRuntime` façade. It keeps provider loops responsive by moving compiler, vault, and telemetry work to an executor, propagates task or token cancellation, and uses native cancellable subprocess paths for RTK command rewriting and aggressive reads. The synchronous `Runtime` API remains unchanged. + It does **not** replace the host's context engine, memory system, transcript store, or provider client. It does not add an MCP server or standing prompt text. If storage, recovery, middleware, or compilation is unavailable or unsafe, the host receives the original request or result unchanged. ## What it does @@ -228,6 +230,39 @@ def reduce_provider_request(request, *, session_id, request_id): An adapter must preserve four contracts: stable request/session identity, original-object immutability, pass-through on `None` or error, and model access to `artifact_get`. The `Runtime` surface is usable today; framework-specific one-command adapters beyond Hermes are not yet shipped. +### Async runtimes, cancellation, and concurrency + +Wrap the same synchronous runtime when the host owns an asyncio event loop: + +```python +from rtk_hermes_plus import AsyncRuntime, CancellationToken, Runtime + +terminator = AsyncRuntime(Runtime(config, profile_name="my-agent")) +cancellation = CancellationToken() + +reduced = await terminator.transform_tool_result( + tool_name="search_files", + args={"pattern": "ContextEngine"}, + result=large_result, + session_id=session_id, + tool_call_id=call_id, + cancellation=cancellation, +) + +compiled = await terminator.llm_request_middleware( + request=provider_request, + session_id=session_id, + request_id=request_id, + cancellation=cancellation, +) +``` + +`AsyncRuntime` mirrors the adapter-facing tool, result, request, recovery, and session methods. Pass a custom `concurrent.futures.Executor` to `AsyncRuntime(runtime, executor=...)` when the host needs a dedicated worker pool. + +Cancelling the awaiting asyncio task, or calling the thread-safe `cancellation.cancel()`, raises `asyncio.CancelledError` at the adapter boundary. Active RTK subprocesses are killed and reaped. Compiler and SQLite operations already running in an executor remain atomic and may finish in that worker after the caller has stopped waiting; their result is discarded. Token Terminator does not intercept or buffer provider response streams, so adapters compile immediately before dispatch and leave streaming responses under host control. + +The artifact vault enables SQLite WAL mode and uses `synchronous=NORMAL`, a ten-second busy timeout, short-lived connections, and `BEGIN IMMEDIATE` writes. Independent runtime instances and agent processes may share one local vault while preserving content deduplication, lease limits, and observation provenance. Keep the database on a local filesystem: SQLite WAL is not a network-filesystem coordination protocol. + ## Exact recovery Compressed results and request receipts contain an artifact identifier. The model can recover an exact page through the registered tool: @@ -302,7 +337,7 @@ A valid comparison requires separate fresh sessions with stable modes, the same ## Security and privacy - Exact raw artifacts and their private provenance are stored locally because recovery is part of the product contract. -- The vault enforces per-artifact and total-capacity limits, SQLite foreign keys, busy timeouts, schema-version checks, and short-lived transactions. +- The vault enforces per-artifact and total-capacity limits, SQLite WAL, foreign keys, busy timeouts, schema-version checks, short-lived transactions, and serialized writes. - POSIX storage uses `0700` parent directories and `0600` databases. Windows storage inherits the user's profile ACLs. - RTK subprocesses use argument arrays with `shell=False`. - Remote terminal backends are disabled by default. diff --git a/src/rtk_hermes_plus/__init__.py b/src/rtk_hermes_plus/__init__.py index de9e747..f6dc307 100644 --- a/src/rtk_hermes_plus/__init__.py +++ b/src/rtk_hermes_plus/__init__.py @@ -6,6 +6,14 @@ """ from ._version import __version__ +from .async_runtime import AsyncRuntime +from .cancellation import CancellationToken from .plugin import Runtime, register -__all__ = ["Runtime", "__version__", "register"] +__all__ = [ + "AsyncRuntime", + "CancellationToken", + "Runtime", + "__version__", + "register", +] diff --git a/src/rtk_hermes_plus/async_runtime.py b/src/rtk_hermes_plus/async_runtime.py new file mode 100644 index 0000000..a0b42df --- /dev/null +++ b/src/rtk_hermes_plus/async_runtime.py @@ -0,0 +1,242 @@ +from __future__ import annotations + +import asyncio +from collections.abc import Callable +from concurrent.futures import Executor +from contextlib import suppress +from functools import partial +from typing import Any, TypeVar + +from .cancellation import CancellationToken +from .compress import compact_text +from .plugin import Runtime + +_T = TypeVar("_T") + + +class AsyncRuntime: + """Non-blocking façade over :class:`Runtime` for async agent frameworks. + + CPU-bound and SQLite-backed sync operations run in an executor. Cancelling + the awaiting task propagates immediately to the caller; already-running + sync work remains atomic and may finish in its worker thread. RTK command + rewrites use a native asyncio subprocess path and are killed and reaped on + cancellation. + """ + + def __init__(self, runtime: Runtime, *, executor: Executor | None = None) -> None: + self.runtime = runtime + self.executor = executor + + def _raise_if_cancelled(self, cancellation: CancellationToken | None) -> None: + if cancellation is not None and cancellation.cancelled: + self.runtime.metrics.add("async_cancelled") + raise asyncio.CancelledError + + async def _run_sync( + self, + function: Callable[..., _T], + /, + *args: Any, + cancellation: CancellationToken | None = None, + **kwargs: Any, + ) -> _T: + self._raise_if_cancelled(cancellation) + loop = asyncio.get_running_loop() + future = loop.run_in_executor(self.executor, partial(function, *args, **kwargs)) + cancellation_waiter: asyncio.Task[None] | None = None + try: + if cancellation is None: + return await future + cancellation_waiter = asyncio.create_task(cancellation.wait()) + done, _pending = await asyncio.wait( + (future, cancellation_waiter), + return_when=asyncio.FIRST_COMPLETED, + ) + if future in done: + return future.result() + if cancellation_waiter in done: + future.cancel() + raise asyncio.CancelledError + raise RuntimeError("async runtime wait completed without a result") + except asyncio.CancelledError: + if cancellation is not None: + cancellation.cancel() + future.cancel() + self.runtime.metrics.add("async_cancelled") + raise + finally: + if cancellation_waiter is not None: + cancellation_waiter.cancel() + with suppress(asyncio.CancelledError): + await cancellation_waiter + + async def tool_request_middleware( + self, + *, + tool_name: str, + args: dict, + cancellation: CancellationToken | None = None, + **kwargs: Any, + ) -> dict | None: + self._raise_if_cancelled(cancellation) + if tool_name != "terminal" or not isinstance(args, dict): + return None + prepared = await self._run_sync( + self.runtime._prepare_rewrite, + args, + cancellation=cancellation, + ) + if prepared is None: + return None + command, cwd = prepared + try: + result = await self.runtime.rewriter.rewrite_async( + command, + cwd=cwd, + cancellation=cancellation, + ) + except asyncio.CancelledError: + if cancellation is not None: + cancellation.cancel() + self.runtime.metrics.add("async_cancelled") + raise + rewritten = self.runtime._apply_rewrite_result(args, result) + if rewritten is None: + return None + await self._run_sync( + self.runtime._record_rewrite, + session_id=str(kwargs.get("session_id") or ""), + turn_id=str(kwargs.get("turn_id") or ""), + cancellation=cancellation, + ) + return { + "args": rewritten, + "source": "token-terminator", + "reason": "strict token reduction", + } + + async def observe_tool_call( + self, *, cancellation: CancellationToken | None = None, **kwargs: Any + ) -> None: + await self._run_sync( + self.runtime.observe_tool_call, cancellation=cancellation, **kwargs + ) + + async def transform_tool_result( + self, + *, + tool_name: str, + args: dict, + result: str, + cancellation: CancellationToken | None = None, + **kwargs: Any, + ) -> Any: + self._raise_if_cancelled(cancellation) + compressor = self.runtime.compressor + if not compressor._eligible(tool_name=tool_name, args=args, result=result): + return None + compressor.metrics.add("native_attempted") + compact = None + try: + if tool_name == "read_file": + compact = await compressor._rtk_read_async( + args, + cancellation=cancellation, + ) + except asyncio.CancelledError: + if cancellation is not None: + cancellation.cancel() + self.runtime.metrics.add("async_cancelled") + raise + if compact is None: + compact = await self._run_sync( + compact_text, + result, + compressor.config.native_max_chars, + cancellation=cancellation, + ) + transformed = await self._run_sync( + compressor._finalize_transform, + tool_name=tool_name, + args=args, + result=result, + compact=compact, + session_id=str(kwargs.get("session_id") or ""), + tool_call_id=str(kwargs.get("tool_call_id") or ""), + cancellation=cancellation, + ) + if transformed is not None: + await self._run_sync( + self.runtime._record_native, + session_id=str(kwargs.get("session_id") or ""), + turn_id=str(kwargs.get("turn_id") or ""), + raw_chars=len(result), + output_chars=len(transformed), + cancellation=cancellation, + ) + return transformed + + async def post_tool_call( + self, *, cancellation: CancellationToken | None = None, **kwargs: Any + ) -> None: + await self._run_sync( + self.runtime.post_tool_call, cancellation=cancellation, **kwargs + ) + + async def llm_request_middleware( + self, *, cancellation: CancellationToken | None = None, **kwargs: Any + ) -> dict | None: + return await self._run_sync( + self.runtime.llm_request_middleware, cancellation=cancellation, **kwargs + ) + + async def on_session_start( + self, + *, + session_id: str, + cancellation: CancellationToken | None = None, + ) -> None: + await self._run_sync( + self.runtime.on_session_start, + session_id=session_id, + cancellation=cancellation, + ) + + async def pre_llm_call( + self, *, cancellation: CancellationToken | None = None, **kwargs: Any + ) -> None: + await self._run_sync( + self.runtime.pre_llm_call, cancellation=cancellation, **kwargs + ) + + async def on_session_end( + self, *, cancellation: CancellationToken | None = None, **kwargs: Any + ) -> None: + await self._run_sync( + self.runtime.on_session_end, cancellation=cancellation, **kwargs + ) + + async def on_session_finalize( + self, + *, + session_id: str, + cancellation: CancellationToken | None = None, + ) -> None: + await self._run_sync( + self.runtime.on_session_finalize, + session_id=session_id, + cancellation=cancellation, + ) + + async def tool( + self, *, cancellation: CancellationToken | None = None, **kwargs: Any + ) -> str: + return await self._run_sync( + self.runtime.tool, cancellation=cancellation, **kwargs + ) + + async def status( + self, *, cancellation: CancellationToken | None = None + ) -> dict[str, Any]: + return await self._run_sync(self.runtime.status, cancellation=cancellation) diff --git a/src/rtk_hermes_plus/cancellation.py b/src/rtk_hermes_plus/cancellation.py new file mode 100644 index 0000000..1020d4a --- /dev/null +++ b/src/rtk_hermes_plus/cancellation.py @@ -0,0 +1,63 @@ +from __future__ import annotations + +import asyncio +import threading + + +class CancellationToken: + """Thread-safe cancellation signal for async runtime adapters. + + Cancelling a token stops cancellable async work such as RTK subprocesses. + SQLite and compiler operations already running in a worker thread remain + atomic and may finish after the awaiting task has been cancelled. + """ + + def __init__(self) -> None: + self._cancelled = False + self._lock = threading.Lock() + self._waiters: set[tuple[asyncio.AbstractEventLoop, asyncio.Future[None]]] = ( + set() + ) + + @property + def cancelled(self) -> bool: + with self._lock: + return self._cancelled + + @staticmethod + def _resolve_waiter(waiter: asyncio.Future[None]) -> None: + if not waiter.done(): + waiter.set_result(None) + + def cancel(self) -> None: + with self._lock: + if self._cancelled: + return + self._cancelled = True + waiters = tuple(self._waiters) + self._waiters.clear() + for loop, waiter in waiters: + try: + loop.call_soon_threadsafe(self._resolve_waiter, waiter) + except RuntimeError: + # The owning loop closed while cancellation crossed threads. + # No active coroutine remains to wake in that loop. + continue + + async def wait(self) -> None: + loop = asyncio.get_running_loop() + waiter = loop.create_future() + entry = (loop, waiter) + with self._lock: + if self._cancelled: + return + self._waiters.add(entry) + try: + await waiter + finally: + with self._lock: + self._waiters.discard(entry) + + def raise_if_cancelled(self) -> None: + if self.cancelled: + raise asyncio.CancelledError diff --git a/src/rtk_hermes_plus/compress.py b/src/rtk_hermes_plus/compress.py index 34e9326..2bea117 100644 --- a/src/rtk_hermes_plus/compress.py +++ b/src/rtk_hermes_plus/compress.py @@ -1,8 +1,11 @@ from __future__ import annotations +import asyncio import re import subprocess +from contextlib import suppress +from .cancellation import CancellationToken from .config import Config from .metrics import Metrics from .rewrite import backend_enabled, command_workdir, terminal_backend @@ -90,25 +93,47 @@ def __init__( def transform( self, *, tool_name: str, args: dict, result: str, **kwargs ) -> str | None: - if not self.config.native_enabled or not isinstance(result, str): + if not self._eligible(tool_name=tool_name, args=args, result=result): return None + + self.metrics.add("native_attempted") + compact = None + if tool_name == "read_file": + compact = self._rtk_read(args) + if compact is None: + compact = compact_text(result, self.config.native_max_chars) + return self._finalize_transform( + tool_name=tool_name, + args=args, + result=result, + compact=compact, + **kwargs, + ) + + def _eligible(self, *, tool_name: str, args: dict, result: str) -> bool: + if not self.config.native_enabled or not isinstance(result, str): + return False allowed = self.BALANCED_TOOLS | ( self.AGGRESSIVE_TOOLS if self.config.aggressive else frozenset() ) if tool_name not in allowed or len(result) < self.config.native_min_chars: - return None + return False backend = terminal_backend(args) if not backend_enabled(backend, self.config): self.metrics.add("native_skipped_backend") - return None + return False + return True - self.metrics.add("native_attempted") - compact = None - if tool_name == "read_file": - compact = self._rtk_read(args) - if compact is None: - compact = compact_text(result, self.config.native_max_chars) + def _finalize_transform( + self, + *, + tool_name: str, + args: dict, + result: str, + compact: str | None, + **kwargs, + ) -> str | None: if not compact or len(compact) >= len(result): self.metrics.add("native_not_smaller") return None @@ -168,6 +193,77 @@ def _rtk_read(self, args: dict) -> str | None: else None ) + async def _rtk_read_async( + self, + args: dict, + *, + cancellation: CancellationToken | None = None, + ) -> str | None: + if cancellation is not None: + cancellation.raise_if_cancelled() + if not self.rtk_path: + return None + path_value = args.get("path") or args.get("file_path") or args.get("filename") + if not isinstance(path_value, str) or not path_value.strip(): + return None + process: asyncio.subprocess.Process | None = None + communicate_task: asyncio.Task[tuple[bytes, bytes]] | None = None + cancellation_waiter: asyncio.Task[None] | None = None + try: + process = await asyncio.create_subprocess_exec( + self.rtk_path, + "read", + path_value, + "-l", + "aggressive", + cwd=str(command_workdir(args)), + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + communicate_task = asyncio.create_task(process.communicate()) + waiters: set[asyncio.Task] = {communicate_task} + if cancellation is not None: + cancellation_waiter = asyncio.create_task(cancellation.wait()) + waiters.add(cancellation_waiter) + done, _pending = await asyncio.wait( + waiters, + timeout=self.config.timeout_ms / 1000, + return_when=asyncio.FIRST_COMPLETED, + ) + if not done: + with suppress(ProcessLookupError): + process.kill() + await communicate_task + self.metrics.add("native_rtk_timeouts") + return None + if cancellation_waiter is not None and cancellation_waiter in done: + raise asyncio.CancelledError + stdout, _stderr = communicate_task.result() + decoded = stdout.decode(errors="replace") + return decoded if process.returncode == 0 and decoded.strip() else None + except asyncio.CancelledError: + if cancellation is not None: + cancellation.cancel() + if process is not None and process.returncode is None: + with suppress(ProcessLookupError): + process.kill() + if communicate_task is not None: + with suppress(asyncio.CancelledError, ProcessLookupError): + await communicate_task + elif process is not None: + with suppress(asyncio.CancelledError, ProcessLookupError): + await process.wait() + self.metrics.add("native_rtk_cancelled") + raise + except OSError: + self.metrics.add("native_rtk_errors") + return None + finally: + if cancellation_waiter is not None: + cancellation_waiter.cancel() + with suppress(asyncio.CancelledError): + await cancellation_waiter + def compact_text(text: str, max_chars: int) -> str: max_chars = max(1, int(max_chars)) diff --git a/src/rtk_hermes_plus/plugin.py b/src/rtk_hermes_plus/plugin.py index 4e0bf19..9f0d97b 100644 --- a/src/rtk_hermes_plus/plugin.py +++ b/src/rtk_hermes_plus/plugin.py @@ -17,6 +17,7 @@ from .metrics import Metrics from .rewrite import ( Rewriter, + RewriteResult, backend_enabled, command_excluded, command_workdir, @@ -102,8 +103,7 @@ def tool_request_middleware(self, *, tool_name: str, args: dict, **kwargs): rewritten = self._rewrite_args(args) if rewritten is None: return None - self._ensure_ledger_session(kwargs.get("session_id")) - self.ledger.record_rewrite( + self._record_rewrite( session_id=str(kwargs.get("session_id") or ""), turn_id=str(kwargs.get("turn_id") or ""), ) @@ -122,8 +122,7 @@ def pre_tool_call(self, *, tool_name: str, args: dict, **kwargs) -> None: if rewritten is not None: args.clear() args.update(rewritten) - self._ensure_ledger_session(kwargs.get("session_id")) - self.ledger.record_rewrite( + self._record_rewrite( session_id=str(kwargs.get("session_id") or ""), turn_id=str(kwargs.get("turn_id") or ""), ) @@ -150,8 +149,7 @@ def transform_tool_result( tool_name=tool_name, args=args, result=result, **kwargs ) if transformed is not None: - self._ensure_ledger_session(kwargs.get("session_id")) - self.ledger.record_native( + self._record_native( session_id=str(kwargs.get("session_id") or ""), turn_id=str(kwargs.get("turn_id") or ""), raw_chars=len(result), @@ -356,10 +354,30 @@ def _ensure_ledger_session(self, session_id) -> None: if session_id: self.ledger.ensure_session(str(session_id), self.config.mode) + def _record_rewrite(self, *, session_id: str = "", turn_id: str = "") -> None: + self._ensure_ledger_session(session_id) + self.ledger.record_rewrite(session_id=session_id, turn_id=turn_id) + + def _record_native( + self, + *, + session_id: str = "", + turn_id: str = "", + raw_chars: int, + output_chars: int, + ) -> None: + self._ensure_ledger_session(session_id) + self.ledger.record_native( + session_id=session_id, + turn_id=turn_id, + raw_chars=raw_chars, + output_chars=output_chars, + ) + # ------------------------------------------------------------------ # Terminal rewrite implementation # ------------------------------------------------------------------ - def _rewrite_args(self, args: dict) -> dict | None: + def _prepare_rewrite(self, args: dict) -> tuple[str, Path] | None: if not self.config.terminal_enabled: return None command = args.get("command") @@ -387,7 +405,9 @@ def _rewrite_args(self, args: dict) -> dict | None: self.metrics.add("rewrite_skipped_pytest_quiet_config") return None - result = self.rewriter.rewrite(command, cwd=cwd) + return command, cwd + + def _apply_rewrite_result(self, args: dict, result: RewriteResult) -> dict | None: if result.command is None: self.metrics.add("rewrite_passthrough") return None @@ -406,6 +426,15 @@ def _rewrite_args(self, args: dict) -> dict | None: self.metrics.add("rewritten") return output + def _rewrite_args(self, args: dict) -> dict | None: + prepared = self._prepare_rewrite(args) + if prepared is None: + return None + command, cwd = prepared + + result = self.rewriter.rewrite(command, cwd=cwd) + return self._apply_rewrite_result(args, result) + # ------------------------------------------------------------------ # One compact model tool for exact recovery and working-state control # ------------------------------------------------------------------ @@ -476,6 +505,7 @@ def status(self) -> dict[str, Any]: "enabled": self.config.enabled, "vault_available": self.store is not None, "vault_error": self.store_error, + "journal_mode": self.store.journal_mode if self.store else "unavailable", "profile": self.profile_name, } diff --git a/src/rtk_hermes_plus/rewrite.py b/src/rtk_hermes_plus/rewrite.py index 2a72fa7..97b223a 100644 --- a/src/rtk_hermes_plus/rewrite.py +++ b/src/rtk_hermes_plus/rewrite.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio import configparser import os import re @@ -8,9 +9,11 @@ import threading import time from collections import OrderedDict +from contextlib import suppress from dataclasses import dataclass from pathlib import Path +from .cancellation import CancellationToken from .config import Config from .metrics import Metrics @@ -75,6 +78,25 @@ def __init__(self, config: Config, metrics: Metrics): def available(self) -> bool: return self.rtk_path is not None + def _result_from_output( + self, + command: str, + *, + stdout: str, + returncode: int, + elapsed_ms: float, + ) -> RewriteResult: + self.metrics.add("rewrite_total_ms", elapsed_ms) + rewritten = stdout.strip() + command_out = ( + rewritten + if returncode in {0, 3} and rewritten and rewritten != command + else None + ) + result = RewriteResult(command_out, returncode, elapsed_ms) + self.cache.put(command, result) + return result + def rewrite(self, command: str, *, cwd: Path) -> RewriteResult: cached = self.cache.get(command) if cached is not None: @@ -94,16 +116,12 @@ def rewrite(self, command: str, *, cwd: Path) -> RewriteResult: check=False, ) elapsed = (time.perf_counter() - started) * 1000 - self.metrics.add("rewrite_total_ms", elapsed) - rewritten = completed.stdout.strip() - command_out = ( - rewritten - if completed.returncode in {0, 3} and rewritten and rewritten != command - else None + return self._result_from_output( + command, + stdout=completed.stdout, + returncode=completed.returncode, + elapsed_ms=elapsed, ) - result = RewriteResult(command_out, completed.returncode, elapsed) - self.cache.put(command, result) - return result except subprocess.TimeoutExpired: elapsed = (time.perf_counter() - started) * 1000 self.metrics.add("rewrite_total_ms", elapsed) @@ -115,6 +133,91 @@ def rewrite(self, command: str, *, cwd: Path) -> RewriteResult: self.metrics.add("rewrite_errors") return RewriteResult(None, -1, elapsed) + async def rewrite_async( + self, + command: str, + *, + cwd: Path, + cancellation: CancellationToken | None = None, + ) -> RewriteResult: + """Rewrite through a cancellable asyncio subprocess.""" + if cancellation is not None: + cancellation.raise_if_cancelled() + cached = self.cache.get(command) + if cached is not None: + self.metrics.add("rewrite_cache_hits") + return cached + + started = time.perf_counter() + self.metrics.add("rewrite_attempted") + process: asyncio.subprocess.Process | None = None + communicate_task: asyncio.Task[tuple[bytes, bytes]] | None = None + cancellation_waiter: asyncio.Task[None] | None = None + try: + process = await asyncio.create_subprocess_exec( + self.rtk_path or "rtk", + "rewrite", + command, + cwd=str(cwd), + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + ) + communicate_task = asyncio.create_task(process.communicate()) + waiters: set[asyncio.Task] = {communicate_task} + if cancellation is not None: + cancellation_waiter = asyncio.create_task(cancellation.wait()) + waiters.add(cancellation_waiter) + done, _pending = await asyncio.wait( + waiters, + timeout=self.config.timeout_ms / 1000, + return_when=asyncio.FIRST_COMPLETED, + ) + if not done: + with suppress(ProcessLookupError): + process.kill() + await communicate_task + elapsed = (time.perf_counter() - started) * 1000 + self.metrics.add("rewrite_total_ms", elapsed) + self.metrics.add("rewrite_timeouts") + return RewriteResult(None, -1, elapsed) + if cancellation_waiter is not None and cancellation_waiter in done: + raise asyncio.CancelledError + + stdout, _stderr = communicate_task.result() + elapsed = (time.perf_counter() - started) * 1000 + return self._result_from_output( + command, + stdout=stdout.decode(errors="replace"), + returncode=int(process.returncode or 0), + elapsed_ms=elapsed, + ) + except asyncio.CancelledError: + if cancellation is not None: + cancellation.cancel() + if process is not None and process.returncode is None: + with suppress(ProcessLookupError): + process.kill() + if communicate_task is not None: + with suppress(asyncio.CancelledError, ProcessLookupError): + await communicate_task + elif process is not None: + with suppress(asyncio.CancelledError, ProcessLookupError): + await process.wait() + elapsed = (time.perf_counter() - started) * 1000 + self.metrics.add("rewrite_total_ms", elapsed) + self.metrics.add("rewrite_cancelled") + raise + except OSError: + elapsed = (time.perf_counter() - started) * 1000 + self.metrics.add("rewrite_total_ms", elapsed) + self.metrics.add("rewrite_errors") + return RewriteResult(None, -1, elapsed) + finally: + if cancellation_waiter is not None: + cancellation_waiter.cancel() + with suppress(asyncio.CancelledError): + await cancellation_waiter + def terminal_backend(args: dict | None = None) -> str: args = args or {} diff --git a/src/rtk_hermes_plus/storage.py b/src/rtk_hermes_plus/storage.py index 877068a..dde0dff 100644 --- a/src/rtk_hermes_plus/storage.py +++ b/src/rtk_hermes_plus/storage.py @@ -5,7 +5,7 @@ import os import sqlite3 from collections.abc import Iterator -from contextlib import contextmanager +from contextlib import closing, contextmanager from dataclasses import dataclass from datetime import datetime, timezone from pathlib import Path @@ -96,15 +96,28 @@ def __init__( self.path.parent.mkdir(parents=True, exist_ok=True, mode=0o700) if os.name == "posix": os.chmod(self.path.parent, 0o700) + self.journal_mode = self._enable_wal() self._initialize() if os.name == "posix": os.chmod(self.path, 0o600) + def _enable_wal(self) -> str: + with closing( + sqlite3.connect(self.path, timeout=10.0, isolation_level=None) + ) as conn: + conn.execute("PRAGMA busy_timeout=10000") + row = conn.execute("PRAGMA journal_mode=WAL").fetchone() + mode = str(row[0] if row else "").lower() + if mode != "wal": + raise RuntimeError(f"artifact vault requires WAL mode, got {mode!r}") + return mode + def _new_connection(self) -> sqlite3.Connection: conn = sqlite3.connect(self.path, timeout=10.0, isolation_level=None) conn.row_factory = sqlite3.Row conn.execute("PRAGMA foreign_keys=ON") conn.execute("PRAGMA busy_timeout=10000") + conn.execute("PRAGMA synchronous=NORMAL") return conn @contextmanager diff --git a/tests/test_async_runtime.py b/tests/test_async_runtime.py new file mode 100644 index 0000000..33bfbaa --- /dev/null +++ b/tests/test_async_runtime.py @@ -0,0 +1,314 @@ +from __future__ import annotations + +import asyncio +import threading + +import pytest + +from rtk_hermes_plus import AsyncRuntime, CancellationToken, Runtime +from rtk_hermes_plus.config import Config +from rtk_hermes_plus.metrics import Metrics +from rtk_hermes_plus.rewrite import Rewriter, RewriteResult + + +def _runtime(tmp_path, **overrides) -> Runtime: + values = { + "mode": "balanced", + "db_path": tmp_path / "artifacts.sqlite3", + "ledger_path": tmp_path / "experiments.sqlite3", + "state_db_path": tmp_path / "state.db", + "ledger_enabled": False, + "context_compaction_enabled": False, + "min_artifact_chars": 20, + "inline_lease_exposures": 0, + } + values.update(overrides) + return Runtime(Config(**values), profile_name="async-test") + + +def test_async_runtime_compiles_without_mutating_caller_request(tmp_path): + runtime = _runtime(tmp_path) + async_runtime = AsyncRuntime(runtime) + request = { + "messages": [ + { + "role": "assistant", + "tool_calls": [ + { + "id": "c1", + "function": {"name": "search_files", "arguments": "{}"}, + } + ], + }, + {"role": "tool", "tool_call_id": "c1", "content": "evidence " * 40}, + ] + } + original = {"messages": [dict(message) for message in request["messages"]]} + + result = asyncio.run( + async_runtime.llm_request_middleware( + request=request, + session_id="session-1", + request_id="request-1", + ) + ) + + assert result is not None + assert result["request"] is not request + assert request == original + assert runtime.store is not None + assert runtime.store.counts()["artifacts"] == 1 + + +def test_async_runtime_rewrites_terminal_request(tmp_path, monkeypatch): + runtime = _runtime(tmp_path, mode="terminal") + async_runtime = AsyncRuntime(runtime) + + async def rewrite_async(command, **_kwargs): + assert command == "git status" + return RewriteResult("rtk git status", 0, 1.0) + + monkeypatch.setattr(runtime.rewriter, "rewrite_async", rewrite_async) + + result = asyncio.run( + async_runtime.tool_request_middleware( + tool_name="terminal", + args={"command": "git status", "cwd": str(tmp_path)}, + session_id="session-1", + turn_id="turn-1", + ) + ) + + assert result is not None + assert result["args"]["command"] == "rtk git status" + assert result["source"] == "token-terminator" + + +def test_async_runtime_compresses_and_vaults_native_result(tmp_path): + runtime = _runtime( + tmp_path, + native_min_chars=20, + native_max_chars=80, + ) + async_runtime = AsyncRuntime(runtime) + original = "repeated evidence line\n" * 100 + + transformed = asyncio.run( + async_runtime.transform_tool_result( + tool_name="search_files", + args={"pattern": "evidence"}, + result=original, + session_id="session-1", + tool_call_id="call-1", + ) + ) + + assert transformed is not None + assert len(transformed) < len(original) + assert "full artifact=" in transformed + assert runtime.store is not None + assert runtime.store.counts()["artifacts"] == 1 + + +def test_async_runtime_offloads_sync_work_from_event_loop(tmp_path, monkeypatch): + runtime = _runtime(tmp_path) + async_runtime = AsyncRuntime(runtime) + started = threading.Event() + release = threading.Event() + + def blocking_status(): + started.set() + release.wait() + return {"ready": True} + + monkeypatch.setattr(runtime, "status", blocking_status) + + async def exercise(): + task = asyncio.create_task(async_runtime.status()) + await asyncio.to_thread(started.wait) + release.set() + return await task + + assert asyncio.run(exercise()) == {"ready": True} + + +def test_cancellation_token_stops_waiting_for_sync_work(tmp_path, monkeypatch): + runtime = _runtime(tmp_path) + async_runtime = AsyncRuntime(runtime) + cancellation = CancellationToken() + started = threading.Event() + release = threading.Event() + + def blocking_status(): + started.set() + release.wait() + return {"late": True} + + monkeypatch.setattr(runtime, "status", blocking_status) + + async def exercise(): + task = asyncio.create_task(async_runtime.status(cancellation=cancellation)) + await asyncio.to_thread(started.wait) + cancellation.cancel() + try: + with pytest.raises(asyncio.CancelledError): + await task + finally: + release.set() + + asyncio.run(exercise()) + assert runtime.metrics.snapshot()["async_cancelled"] == 1 + + +def test_pre_cancelled_token_does_not_start_work(tmp_path, monkeypatch): + runtime = _runtime(tmp_path) + async_runtime = AsyncRuntime(runtime) + cancellation = CancellationToken() + cancellation.cancel() + called = False + + def status(): + nonlocal called + called = True + return {} + + monkeypatch.setattr(runtime, "status", status) + + async def exercise(): + with pytest.raises(asyncio.CancelledError): + await async_runtime.status(cancellation=cancellation) + + asyncio.run(exercise()) + assert called is False + assert runtime.metrics.snapshot()["async_cancelled"] == 1 + + +def test_cancellation_token_can_be_triggered_from_another_thread(): + cancellation = CancellationToken() + + async def exercise(): + waiter = asyncio.create_task(cancellation.wait()) + thread = threading.Thread(target=cancellation.cancel) + thread.start() + await waiter + thread.join() + + asyncio.run(exercise()) + assert cancellation.cancelled is True + + +class _FakeProcess: + def __init__(self) -> None: + self.returncode = None + self.started = asyncio.Event() + self.finished = asyncio.Event() + self.killed = False + + async def communicate(self): + self.started.set() + await self.finished.wait() + return b"rtk git status", b"" + + def kill(self) -> None: + self.killed = True + self.returncode = -9 + self.finished.set() + + async def wait(self) -> int: + await self.finished.wait() + return int(self.returncode or 0) + + +def test_async_rewriter_kills_and_reaps_on_token_cancellation(tmp_path, monkeypatch): + rewriter = Rewriter( + Config(mode="terminal", timeout_ms=10_000), + Metrics(), + ) + rewriter.rtk_path = "rtk" + process = _FakeProcess() + + async def create_subprocess_exec(*_args, **_kwargs): + return process + + monkeypatch.setattr(asyncio, "create_subprocess_exec", create_subprocess_exec) + + async def exercise(): + cancellation = CancellationToken() + task = asyncio.create_task( + rewriter.rewrite_async( + "git status", + cwd=tmp_path, + cancellation=cancellation, + ) + ) + await process.started.wait() + cancellation.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + asyncio.run(exercise()) + assert process.killed is True + assert rewriter.metrics.snapshot()["rewrite_cancelled"] == 1 + + +def test_async_rewriter_kills_and_reaps_on_task_cancellation(tmp_path, monkeypatch): + rewriter = Rewriter( + Config(mode="terminal", timeout_ms=10_000), + Metrics(), + ) + rewriter.rtk_path = "rtk" + process = _FakeProcess() + + async def create_subprocess_exec(*_args, **_kwargs): + return process + + monkeypatch.setattr(asyncio, "create_subprocess_exec", create_subprocess_exec) + + async def exercise(): + task = asyncio.create_task(rewriter.rewrite_async("git status", cwd=tmp_path)) + await process.started.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + asyncio.run(exercise()) + assert process.killed is True + assert rewriter.metrics.snapshot()["rewrite_cancelled"] == 1 + + +def test_async_native_read_kills_rtk_before_vault_write(tmp_path, monkeypatch): + runtime = _runtime( + tmp_path, + mode="aggressive", + native_min_chars=20, + ) + runtime.compressor.rtk_path = "rtk" + async_runtime = AsyncRuntime(runtime) + process = _FakeProcess() + + async def create_subprocess_exec(*_args, **_kwargs): + return process + + monkeypatch.setattr(asyncio, "create_subprocess_exec", create_subprocess_exec) + + async def exercise(): + cancellation = CancellationToken() + task = asyncio.create_task( + async_runtime.transform_tool_result( + tool_name="read_file", + args={"path": "large.py"}, + result="evidence " * 100, + cancellation=cancellation, + ) + ) + await process.started.wait() + cancellation.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + asyncio.run(exercise()) + assert process.killed is True + assert runtime.metrics.snapshot()["native_rtk_cancelled"] == 1 + assert runtime.metrics.snapshot()["async_cancelled"] == 1 + assert runtime.store is not None + assert runtime.store.counts()["artifacts"] == 0 diff --git a/tests/test_storage_concurrency.py b/tests/test_storage_concurrency.py new file mode 100644 index 0000000..37e1c89 --- /dev/null +++ b/tests/test_storage_concurrency.py @@ -0,0 +1,113 @@ +from __future__ import annotations + +import multiprocessing +from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor +from pathlib import Path + +from rtk_hermes_plus.storage import TokenTerminatorStore + + +def _write_artifacts(path: str, worker: int, iterations: int, variants: int) -> int: + store = TokenTerminatorStore(path) + for index in range(iterations): + store.put_artifact( + f"shared artifact {index % variants}", + tool_name="stress_tool", + args={"variant": index % variants}, + session_id=f"worker-{worker}", + tool_call_id=f"call-{worker}-{index}", + ) + return iterations + + +def _claim_exposure(path: str, artifact_id: str, request_id: str) -> bool: + store = TokenTerminatorStore(path) + return store.claim_exposure( + session_id="shared-session", + artifact_id=artifact_id, + request_id=request_id, + inline_limit=1, + ) + + +def test_artifact_vault_uses_wal_mode(tmp_path): + store = TokenTerminatorStore(tmp_path / "wal.db") + + assert store.journal_mode == "wal" + with store.connection() as connection: + assert connection.execute("PRAGMA journal_mode").fetchone()[0] == "wal" + assert connection.execute("PRAGMA synchronous").fetchone()[0] == 1 + + +def test_concurrent_agent_threads_preserve_dedup_and_observations(tmp_path): + path = str(tmp_path / "threads.db") + workers = 8 + iterations = 40 + variants = 7 + + with ThreadPoolExecutor(max_workers=workers) as executor: + completed = list( + executor.map( + _write_artifacts, + [path] * workers, + range(workers), + [iterations] * workers, + [variants] * workers, + ) + ) + + store = TokenTerminatorStore(path) + counts = store.counts() + assert completed == [iterations] * workers + assert counts["artifacts"] == variants + assert counts["artifact_observations"] == workers * iterations + with store.connection() as connection: + assert connection.execute("PRAGMA integrity_check").fetchone()[0] == "ok" + + +def test_concurrent_agent_processes_preserve_dedup_and_observations(tmp_path): + path = str(Path(tmp_path) / "processes.db") + workers = 4 + iterations = 25 + variants = 5 + context = multiprocessing.get_context("spawn") + + with ProcessPoolExecutor(max_workers=workers, mp_context=context) as executor: + futures = [ + executor.submit(_write_artifacts, path, worker, iterations, variants) + for worker in range(workers) + ] + completed = [future.result() for future in futures] + + store = TokenTerminatorStore(path) + counts = store.counts() + assert completed == [iterations] * workers + assert counts["artifacts"] == variants + assert counts["artifact_observations"] == workers * iterations + with store.connection() as connection: + assert connection.execute("PRAGMA integrity_check").fetchone()[0] == "ok" + + +def test_concurrent_agent_processes_claim_only_one_inline_lease(tmp_path): + path = str(Path(tmp_path) / "leases.db") + store = TokenTerminatorStore(path) + artifact_id = store.put_artifact("shared evidence").artifact_id + workers = 6 + context = multiprocessing.get_context("spawn") + + with ProcessPoolExecutor(max_workers=workers, mp_context=context) as executor: + futures = [ + executor.submit(_claim_exposure, path, artifact_id, f"request-{index}") + for index in range(workers) + ] + decisions = [future.result() for future in futures] + + assert sum(decisions) == 1 + with store.connection() as connection: + assert ( + connection.execute( + "SELECT COUNT(*) FROM artifact_exposures WHERE inline=1" + ).fetchone()[0] + == 1 + ) + assert connection.execute("PRAGMA integrity_check").fetchone()[0] == "ok"