diff --git a/README.md b/README.md index bb81891..c46bb1b 100644 --- a/README.md +++ b/README.md @@ -26,8 +26,17 @@ kw.record(turn, answer) # Khwan persist # turn's context — only under load, which makes it read as flaky memory rather # than as a race. Skip the wait when the turn is the last one: kw.record(turn, answer, background=True) # → {"queued": True} + +# The send runs on a daemon thread, and the interpreter does not wait for those. +# In a CLI, a serverless handler, or any script that ends soon after its last +# turn, that write can be killed mid-flight — no error, the turn simply never +# learned. Wait for it before you exit: +kw.flush() # → how many were in flight ``` +A `flush()` also runs automatically at interpreter exit, bounded to five seconds, +so forgetting the call costs latency rather than the turn. + ## On an event loop Every agent framework worth integrating is async, and a blocking client on an diff --git a/pyproject.toml b/pyproject.toml index 65683a6..328cdc4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "khwan" -version = "0.3.0" +version = "0.3.1" description = "Khwan hosted client — the cognition layer (memory + identity + learning) for your own agent. Bring your own model." readme = "README.md" requires-python = ">=3.9" diff --git a/src/khwan/__init__.py b/src/khwan/__init__.py index 72fbadf..002a8c8 100644 --- a/src/khwan/__init__.py +++ b/src/khwan/__init__.py @@ -25,6 +25,7 @@ from __future__ import annotations import asyncio +import atexit import threading from typing import Any, Callable, Dict, List, Optional @@ -34,7 +35,7 @@ import requests -__version__ = "0.3.0" +__version__ = "0.3.1" DEFAULT_BASE_URL = "https://api.khwan.ai" @@ -179,6 +180,14 @@ def __init__(self, *, user_id: Optional[str] = None, api_key: Optional[str] = No # Auto-retry transient failures (429/502/503/504 honoring Retry-After, plus # network errors on idempotent calls) with exponential backoff + jitter. self._max_retries = max(0, max_retries) + # Live background records. A daemon thread is killed the moment the + # interpreter exits, so a fire-and-forget write in a CLI, a serverless + # handler or any short-lived process can be dropped mid-flight with no + # error anywhere — the turn is simply never learned. Tracking them is what + # makes flush() possible, and atexit makes it happen even unasked. + self._pending: "set[threading.Thread]" = set() + self._pending_lock = threading.Lock() + atexit.register(self._flush_at_exit) # session config forwarded to the server (model may be overridden by the # account's dashboard settings; constitution is a named profile reference). self._cfg = {k: v for k, v in @@ -247,10 +256,49 @@ def _send() -> None: {"turn_token": turn.turn_token, "answer": answer}) except Exception: # noqa: BLE001 — a failed learn must not raise into a thread pass - - threading.Thread(target=_send, daemon=True, name="khwan-record").start() + finally: + with self._pending_lock: + self._pending.discard(threading.current_thread()) + + t = threading.Thread(target=_send, daemon=True, name="khwan-record") + with self._pending_lock: + self._pending.add(t) + t.start() return {"queued": True} + def flush(self, timeout: Optional[float] = None) -> int: + """Wait for background records to finish. Returns how many were still in + flight when called. + + Only matters after ``record(background=True)``. Call it before a + short-lived process ends — a CLI, a serverless handler, a script — because + the sending threads are daemons and the interpreter will not wait for them. + + ``timeout`` is the total budget in seconds, not per thread. It returns + rather than raising when the budget runs out: a record that did not land + costs one turn of learning, and turning that into an exception at exit + would be worse than the thing it reports. + """ + with self._pending_lock: + threads = list(self._pending) + deadline = None if timeout is None else time.monotonic() + timeout + for t in threads: + remaining = None if deadline is None else max(0.0, deadline - time.monotonic()) + t.join(remaining) + return len(threads) + + def _flush_at_exit(self) -> None: + """Last-chance flush, bounded so a hung request cannot wedge the exit. + + A developer who never calls flush() should still not silently lose the + last turn of a script, which is the common case and the one that reads as + "the memory is unreliable" rather than as a missing call. + """ + try: + self.flush(timeout=5.0) + except Exception: # noqa: BLE001 — never raise from an interpreter shutdown + pass + def verify(self, turn: Turn, draft: str) -> dict: """Score a draft answer against the brain BEFORE you ship it. diff --git a/test_flush.py b/test_flush.py new file mode 100644 index 0000000..938106c --- /dev/null +++ b/test_flush.py @@ -0,0 +1,128 @@ +"""flush() — the fire-and-forget write that must not vanish at exit. + +record(background=True) sends on a daemon thread, and the interpreter does not +wait for daemons. In a CLI, a serverless handler, or any script that ends soon +after its last turn, the write can be killed mid-flight with no error anywhere — +the turn is simply never learned, and that reads as unreliable memory rather than +as a missing call. + +The real proof is the subprocess test: a process that exits immediately after a +background record still lands the write. + +Run: python3 test_flush.py +""" + +import subprocess +import sys +import threading +import time +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).parent / "src")) + +import requests # noqa: E402 + +from khwan import Khwan, Turn # noqa: E402 + + +class _Resp: + status_code, headers, content = 200, {}, b"{}" + text = "{}" + + def json(self): + return {} + + +def test_flush_waits_for_the_write(): + started, finished = threading.Event(), [] + + def slow(method, url, **kw): + started.set() + time.sleep(0.25) + finished.append(url) + return _Resp() + + real, requests.request = requests.request, slow + try: + kw = Khwan(api_key="k", base_url="https://example.invalid") + kw.record(Turn({"turn_token": "t"}), "a", background=True) + started.wait(2) + assert finished == [], "precondition: the write should still be in flight" + pending = kw.flush() + assert pending == 1, pending + assert len(finished) == 1, "flush returned before the write landed" + finally: + requests.request = real + print("✓ flush: returns only once the in-flight write has landed") + + +def test_flush_is_bounded_and_does_not_raise(): + """A hung request must not turn into a hang at exit, or an exception there.""" + def hang(method, url, **kw): + time.sleep(10) + return _Resp() + + real, requests.request = requests.request, hang + try: + kw = Khwan(api_key="k", base_url="https://example.invalid") + kw.record(Turn({"turn_token": "t"}), "a", background=True) + t0 = time.monotonic() + kw.flush(timeout=0.2) # returns, does not raise + assert time.monotonic() - t0 < 2, "flush ignored its timeout" + finally: + requests.request = real + print("✓ flush: a hung write is given up on, not raised and not waited out") + + +def test_flush_with_nothing_pending_is_a_no_op(): + kw = Khwan(api_key="k", base_url="https://example.invalid") + assert kw.flush() == 0 + print("✓ flush: nothing in flight → nothing to wait for") + + +def test_a_process_that_exits_immediately_still_lands_the_write(): + """The actual failure this fixes, in a real interpreter that really exits. + + Without the atexit hook the daemon thread is killed and `landed` stays empty. + """ + script = r''' +import sys, threading, time +sys.path.insert(0, %r) +import requests +from khwan import Khwan, Turn + +landed = [] + +class R: + status_code, headers, content, text = 200, {}, b"{}", "{}" + def json(self): return {} + +def slow(method, url, **kw): + time.sleep(0.3) # still in flight when the script falls off the end + landed.append(url) + with open(%r, "a") as f: + f.write(url + "\n") + return R() + +requests.request = slow +kw = Khwan(api_key="k", base_url="https://example.invalid") +kw.record(Turn({"turn_token": "t"}), "answer", background=True) +# no flush() call, no sleep — the process ends here, as a CLI would +''' % (str(Path(__file__).parent / "src"), str(Path(__file__).parent / ".flush-probe")) + + probe = Path(__file__).parent / ".flush-probe" + probe.unlink(missing_ok=True) + subprocess.run([sys.executable, "-c", script], check=True, timeout=30) + assert probe.exists() and "/record" in probe.read_text(), ( + "the process exited before the background write landed — the failure " + "flush() exists to prevent") + probe.unlink() + print("✓ exit: a script that ends immediately still lands its last write") + + +if __name__ == "__main__": + test_flush_waits_for_the_write() + test_flush_is_bounded_and_does_not_raise() + test_flush_with_nothing_pending_is_a_no_op() + test_a_process_that_exits_immediately_still_lands_the_write() + print("\n✅ flush verified")