diff --git a/dikte/assistant.py b/dikte/assistant.py index b57fd06..21ff241 100644 --- a/dikte/assistant.py +++ b/dikte/assistant.py @@ -24,13 +24,17 @@ while they work. import json import os +import re import shutil +import signal import subprocess +import tempfile import threading import time from . import api from . import config as cfg +from . import paths from .i18n import t SESSION_FILE = cfg.DATA_DIR / "assistant.json" @@ -114,13 +118,19 @@ def display_name(conf): # one along costs tokens and invites an answer to the wrong question. Switching # provider drops it too, since none of them can pick up another's thread. -def _read_row(name, max_age_seconds): +def _read_session(): + """The stored conversation row, or {} however the file fails to read.""" try: with open(SESSION_FILE, encoding="utf-8") as fh: row = json.load(fh) except (OSError, json.JSONDecodeError, ValueError): return {} - if not isinstance(row, dict) or row.get("provider") != name: + return row if isinstance(row, dict) else {} + + +def _read_row(name, max_age_seconds): + row = _read_session() + if row.get("provider") != name: return {} if max_age_seconds and time.time() - row.get("ts", 0) > max_age_seconds: return {} @@ -159,22 +169,13 @@ def clear_session(): def stored_provider(): """Whose conversation is on disk, whatever the setting says now.""" - try: - with open(SESSION_FILE, encoding="utf-8") as fh: - row = json.load(fh) - except (OSError, json.JSONDecodeError, ValueError): - return "" - return str(row.get("provider", "")) if isinstance(row, dict) else "" + return str(_read_session().get("provider", "")) def session_age(): """Seconds since the stored conversation was last used, or None.""" - try: - with open(SESSION_FILE, encoding="utf-8") as fh: - row = json.load(fh) - except (OSError, json.JSONDecodeError, ValueError): - return None - if not isinstance(row, dict) or not (row.get("session") or row.get("messages")): + row = _read_session() + if not (row.get("session") or row.get("messages")): return None return time.time() - row.get("ts", 0) @@ -373,14 +374,24 @@ def _stream(cmd, conf, on_event, should_stop): Returns (exit code, stderr). Raises Cancelled when the stop was asked for, and AssistantError when the clock ran out. """ + # stderr lands in a file rather than a pipe: nobody drains it while stdout + # is being read, and a CLI chatty enough on stderr would fill the pipe's + # buffer and wedge both of us. A file has no such limit, and is read once + # at the end, which is the only moment stderr matters. + stderr_file = tempfile.TemporaryFile() + # On POSIX the run gets its own session, so that ending it can take down + # every subprocess it started, not just the CLI itself. + grouped = {"start_new_session": True} if os.name == "posix" else {} try: proc = subprocess.Popen( cmd, cwd=working_dir(conf), stdin=subprocess.DEVNULL, - stdout=subprocess.PIPE, stderr=subprocess.PIPE, + stdout=subprocess.PIPE, stderr=stderr_file, text=True, encoding="utf-8", errors="replace", bufsize=1, - creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0), + creationflags=paths.NO_WINDOW, + **grouped, ) except OSError as exc: + stderr_file.close() raise AssistantError(t("Could not run {binary}: {error}", binary=cmd[0], error=exc)) from exc @@ -409,7 +420,7 @@ def _stream(cmd, conf, on_event, should_stop): if isinstance(event, dict): on_event(event) finally: - stderr = _finish(proc) + stderr = _finish(proc, stderr_file) watchdog.join(timeout=1) if ended["cancelled"]: @@ -420,12 +431,31 @@ def _stream(cmd, conf, on_event, should_stop): return proc.returncode, stderr +# Failures that a fresh session cannot cure: an exhausted quota, a signed-out +# CLI, a network that is down. A resumed run that dies with one of these is +# reported as what it is, not retried without the session, because the retry +# would fail the same way after making the user wait through a second run. +_API_TROUBLE = re.compile( + r"(?i)rate.?limit|quota|overloaded|too many requests|credit|billing|" + r"insufficient|unauthorized|forbidden|authentication|invalid.{0,8}key|" + r"log ?in|logged.?out|network|connection|ECONN|ENOTFOUND|ETIMEDOUT|" + r"\b(401|403|429|5\d\d)\b") + + def _conclude(found, code, stderr, session, service): """Turn what the stream said into an answer, or into the reason there is none.""" if code != 0 and not found["answer"]: - if session and _session_missing(stderr): + # A resumed run that died with nothing to show is treated as the + # session being gone, whatever the wording: this code used to look for + # "session ... not found" in stderr, but a CLI update or another + # language rewords that and the recovery stops working. Retrying costs + # one clean start, and cannot loop because the retry resumes nothing. + # Recognised API trouble is the exception: it is not the session's + # fault, and the retry would only repeat it. + blame = last_line(stderr) or found["failure"] or "" + if session and not _API_TROUBLE.search(blame): raise _SessionGone() - raise AssistantError(last_line(stderr) or found["failure"] or t( + raise AssistantError(blame or t( "{service} exited with code {code}.", service=service, code=code)) if found["failure"] and not found["answer"]: raise AssistantError(found["failure"]) @@ -436,14 +466,6 @@ def _conclude(found, code, stderr, session, service): return found["answer"], found["warning"] -def _session_missing(stderr): - lowered = (stderr or "").lower() - if "session" in lowered or "thread" in lowered or "conversation" in lowered: - return any(word in lowered for word in ("not found", "no such", "unknown", - "does not exist", "no conversation")) - return False - - def _watch(proc, deadline, should_stop, ended): while proc.poll() is None: if should_stop is not None and should_stop(): @@ -454,29 +476,60 @@ def _watch(proc, deadline, should_stop, ended): break time.sleep(0.25) if ended["cancelled"] or ended["timed_out"]: - _kill(proc) + kill_tree(proc) -def _kill(proc): +def kill_tree(proc): + """End the process and everything it started. + + A CLI runs tools as subprocesses of its own, and ending only the CLI would + leave those behind, still working on a question nobody is waiting for. + Shared with cleanup, which runs the same two programs. Every failure here + is swallowed: the process being already gone is the outcome being asked for. + """ + if os.name == "nt": + # There is no process group to signal on Windows; taskkill walks the + # tree instead. The wait after it is best-effort, so a tree that will + # not die does not hang the caller on top of everything else. + subprocess.run( + ["taskkill", "/T", "/F", "/PID", str(proc.pid)], + capture_output=True, + creationflags=paths.NO_WINDOW, + ) + try: + proc.wait(timeout=3) + except (subprocess.TimeoutExpired, OSError): + pass + return + # The Popen was started with start_new_session=True, so the pid names a + # whole session to signal. SIGTERM first for a clean exit, SIGKILL for a + # tree that ignored it. + try: + os.killpg(proc.pid, signal.SIGTERM) + except (ProcessLookupError, PermissionError, OSError): + return try: - proc.terminate() proc.wait(timeout=3) except subprocess.TimeoutExpired: - proc.kill() - except OSError: - pass + try: + os.killpg(proc.pid, signal.SIGKILL) + except (ProcessLookupError, PermissionError, OSError): + pass -def _finish(proc): - try: - stderr = proc.stderr.read() or "" - except (OSError, ValueError): - stderr = "" +def _finish(proc, stderr_file): try: proc.wait(timeout=5) except subprocess.TimeoutExpired: - proc.kill() - for stream in (proc.stdout, proc.stderr): + kill_tree(proc) + # Read back what the CLI wrote to its stderr file, decoded leniently: a + # dying CLI is exactly the one likely to print something half-encoded. + try: + stderr_file.seek(0) + stderr = stderr_file.read().decode("utf-8", "replace") + except (OSError, ValueError): + stderr = "" + for stream in (proc.stdout, stderr_file): try: stream.close() except OSError: diff --git a/dikte/cleanup.py b/dikte/cleanup.py index 96eab1d..14e44df 100644 --- a/dikte/cleanup.py +++ b/dikte/cleanup.py @@ -21,6 +21,7 @@ import tempfile from . import api from . import assistant from . import ggml +from . import paths from .i18n import t PROVIDERS = ("openrouter", "local", "claude", "codex") @@ -194,21 +195,43 @@ def _output(cmd, timeout, service): "{binary} not found. Install it, or have OpenRouter clean up " "instead, under Settings → API and models.", binary=binary, )) + # Both streams land in files rather than pipes: nobody drains a pipe while + # the process is being waited out, and a CLI chatty enough would fill the + # buffer and wedge. And a timeout must end the CLI's tool subprocesses too, + # not just the CLI, which subprocess.run's timeout does not do; hence the + # own session on POSIX and assistant.kill_tree on the way out. + out_file = tempfile.TemporaryFile() + err_file = tempfile.TemporaryFile() + grouped = {"start_new_session": True} if os.name == "posix" else {} try: - done = subprocess.run( - cmd, cwd=os.path.expanduser("~"), stdin=subprocess.DEVNULL, - capture_output=True, text=True, encoding="utf-8", errors="replace", - timeout=timeout, - creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0), - ) - except subprocess.TimeoutExpired: - raise CleanupError(t("{service} did not finish within {seconds} seconds.", - service=service, seconds=timeout)) from None - except OSError as exc: - raise CleanupError(t("Could not run {binary}: {error}", - binary=binary, error=exc)) from exc - if done.returncode != 0: - raise CleanupError(assistant.last_line(done.stderr) or t( + try: + proc = subprocess.Popen( + cmd, cwd=os.path.expanduser("~"), stdin=subprocess.DEVNULL, + stdout=out_file, stderr=err_file, + creationflags=paths.NO_WINDOW, + **grouped, + ) + except OSError as exc: + raise CleanupError(t("Could not run {binary}: {error}", + binary=binary, error=exc)) from exc + try: + proc.wait(timeout=timeout) + except subprocess.TimeoutExpired: + assistant.kill_tree(proc) + raise CleanupError(t("{service} did not finish within {seconds} seconds.", + service=service, seconds=timeout)) from None + out_file.seek(0) + stdout = out_file.read().decode("utf-8", "replace") + err_file.seek(0) + stderr = err_file.read().decode("utf-8", "replace") + finally: + for handle in (out_file, err_file): + try: + handle.close() + except OSError: + pass + if proc.returncode != 0: + raise CleanupError(assistant.last_line(stderr) or t( "{service} exited with code {code}.", - service=service, code=done.returncode)) - return (done.stdout or "").strip() + service=service, code=proc.returncode)) + return stdout.strip() diff --git a/tests/test_assistant.py b/tests/test_assistant.py index 1c682c6..faf6362 100644 --- a/tests/test_assistant.py +++ b/tests/test_assistant.py @@ -10,6 +10,8 @@ import io import json import os import subprocess +import sys +import threading import time import unittest from unittest import mock @@ -19,12 +21,16 @@ from tests.support import DikteTest, fake_urlopen, only_these_tools class FakeCli: - """A CLI that prints the given events and exits.""" + """A CLI that prints the given events and exits. - def __init__(self, events=(), code=0, stderr="", noise=()): + Its stderr is not modelled: _stream hands the process a temporary file for + that, and a mocked Popen leaves the file empty, which is what a quiet CLI + writes anyway. + """ + + def __init__(self, events=(), code=0, noise=()): lines = list(noise) + [json.dumps(event) for event in events] self.stdout = io.StringIO("\n".join(lines) + "\n") - self.stderr = io.StringIO(stderr) self.returncode = code self.killed = False @@ -41,6 +47,38 @@ class FakeCli: self.killed = True +class WedgedCli: + """A CLI whose stdout produces nothing, the way a hung process's does. + + Iterating its stdout blocks until the kill arrives, because that is what + reading a silent pipe does: only the process ending closes the stream. + """ + + def __init__(self): + self.pid = 4242 + self.returncode = None + self.released = threading.Event() + self.stdout = self + + def __iter__(self): + return self + + def __next__(self): + # The 5 second cap is a safety net for the test itself; the kill is + # what is supposed to end the wait. + self.released.wait(timeout=5) + raise StopIteration + + def poll(self): + return self.returncode + + def wait(self, timeout=None): + return self.returncode + + def close(self): + pass + + class Provider(DikteTest): def test_the_default(self): self.assertEqual(assistant.provider(self.config()), "claude") @@ -221,19 +259,7 @@ class Denials(DikteTest): self.assertIn("Write", warning) -class SessionMissing(unittest.TestCase): - def test_a_session_that_is_gone(self): - for text in ("Error: session abc not found", - "No conversation with that id", - "unknown thread: abc"): - with self.subTest(text=text): - self.assertTrue(assistant._session_missing(text)) - - def test_an_unrelated_failure(self): - for text in ("", "network unreachable", "session limit exceeded"): - with self.subTest(text=text): - self.assertFalse(assistant._session_missing(text)) - +class LastLine(unittest.TestCase): def test_the_last_line_is_the_one_worth_showing(self): self.assertEqual(assistant.last_line("warning\n\nreal error\n"), "real error") @@ -269,6 +295,28 @@ class Conclude(DikteTest): assistant._conclude(self.found(), 1, "session abc not found", "abc", "Claude") + def test_the_recovery_no_longer_hangs_on_the_words_the_cli_chose(self): + # The complaint used to be matched by substring, which a CLI update or + # another language broke. A resumed run that died with nothing to show + # is now enough on its own. + for stderr in ("Oturum bulunamadı", "something else entirely", ""): + with self.subTest(stderr=stderr): + with self.assertRaises(assistant._SessionGone): + assistant._conclude(self.found(), 1, stderr, "abc", "Claude") + + def test_api_trouble_on_a_resumed_run_is_not_blamed_on_the_session(self): + # A fresh session cannot cure a spent quota, a signed-out CLI or a dead + # network: the retry would fail the same way after a second wait, and + # the user would lose the conversation thread on top. + for stderr in ("Rate limit exceeded", + "You are not logged in. Please run /login.", + "API Error: 401 Unauthorized", + "fetch failed: ECONNREFUSED 127.0.0.1"): + with self.subTest(stderr=stderr): + with self.assertRaises(assistant.AssistantError) as caught: + assistant._conclude(self.found(), 1, stderr, "abc", "Claude") + self.assertIn(stderr, str(caught.exception)) + def test_a_session_that_is_gone_only_matters_when_one_was_resumed(self): with self.assertRaises(assistant.AssistantError): assistant._conclude(self.found(), 1, "session abc not found", @@ -279,6 +327,11 @@ class Conclude(DikteTest): "", "Claude") self.assertEqual(answer, "done") + def test_an_answer_on_a_resumed_session_is_kept_rather_than_retried(self): + answer, _ = assistant._conclude(self.found(answer="done"), 1, "noise", + "abc", "Claude") + self.assertEqual(answer, "done") + def test_a_reported_failure_with_no_answer(self): with self.assertRaises(assistant.AssistantError) as caught: assistant._conclude(self.found(failure="the model refused"), 0, "", @@ -291,14 +344,53 @@ class Conclude(DikteTest): self.assertIn("Codex", str(caught.exception)) +class Stream(DikteTest): + def test_a_cli_that_floods_stderr_still_finishes(self): + # A real subprocess, because the wedge being tested is real plumbing: + # with stderr on a pipe nobody drains, 200 KB fills the pipe's buffer, + # the child blocks writing it, and the run hangs until the watchdog + # timeout. With stderr on a file the run completes at once. + script = ( + "import sys\n" + "sys.stderr.write('x' * 200000)\n" + "sys.stderr.flush()\n" + "print('{\"type\": \"result\", \"result\": \"done\"}')\n" + ) + conf = self.config(assistant_timeout=15) + events = [] + code, stderr = assistant._stream( + [sys.executable, "-c", script], conf, events.append, None) + self.assertEqual(code, 0) + self.assertEqual(len(stderr), 200000) + self.assertEqual(events[-1]["result"], "done") + + def test_the_watchdog_takes_the_whole_tree_down_on_timeout(self): + proc = WedgedCli() + + def killed(target): + # What the real kill does, as far as _stream can see: the process + # ends, and its closing stream releases the blocked read. + target.returncode = 1 + target.released.set() + + conf = self.config(assistant_timeout=0) + with mock.patch.object(subprocess, "Popen", return_value=proc), \ + mock.patch.object(assistant, "kill_tree", + side_effect=killed) as kill: + with self.assertRaises(assistant.AssistantError) as caught: + assistant._stream(["claude"], conf, lambda event: None, None) + kill.assert_called_once_with(proc) + self.assertIn("did not finish", str(caught.exception)) + + class AskClaude(DikteTest): - def run_ask(self, conf=None, events=None, code=0, stderr="", noise=(), + def run_ask(self, conf=None, events=None, code=0, noise=(), session=""): conf = conf or self.config() proc = FakeCli(events or [ {"type": "system", "subtype": "init", "session_id": "abc"}, {"type": "result", "session_id": "abc", "result": " done "}, - ], code=code, stderr=stderr, noise=noise) + ], code=code, noise=noise) stages = [] with only_these_tools("claude", "codex"), \ mock.patch.object(subprocess, "Popen", return_value=proc) as popen: @@ -533,6 +625,44 @@ class Ask(DikteTest): self.assertEqual(attempts, ["stale-id", ""]) self.assertEqual(assistant.stored_provider(), "") + def test_a_resumed_run_that_dies_is_retried_without_the_session_flag(self): + # All the way through the stream this time: the first run exits 1 with + # no answer and whatever stderr it liked, and the recovery must not + # depend on those words. + conf = self.config() + assistant.write_session("claude", "stale-id") + procs = iter([ + FakeCli(code=1), + FakeCli(events=[{"type": "result", "result": "done"}]), + ]) + cmds = [] + + def popen(cmd, **kwargs): + cmds.append(cmd) + return next(procs) + + with only_these_tools("claude"), \ + mock.patch.object(subprocess, "Popen", side_effect=popen): + answer, _ = assistant.ask("hi", conf) + self.assertEqual(answer, "done") + self.assertEqual(cmds[0][cmds[0].index("--resume") + 1], "stale-id") + self.assertNotIn("--resume", cmds[1]) + + def test_a_run_that_dies_with_an_answer_in_hand_is_not_retried(self): + conf = self.config() + assistant.write_session("claude", "stale-id") + calls = [] + + def popen(cmd, **kwargs): + calls.append(cmd) + return FakeCli(events=[{"type": "result", "result": "done"}], code=1) + + with only_these_tools("claude"), \ + mock.patch.object(subprocess, "Popen", side_effect=popen): + answer, _ = assistant.ask("hi", conf) + self.assertEqual(answer, "done") + self.assertEqual(len(calls), 1) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_cleanup.py b/tests/test_cleanup.py index 7753029..e4222a5 100644 --- a/tests/test_cleanup.py +++ b/tests/test_cleanup.py @@ -1,9 +1,9 @@ -"""Who cleans the transcript up, and what they are asked. +"""Who cleans the transcript up, and what they are asked. -The CLIs are faked at subprocess.run: what the tests read is the argument list -each one is given, where the answer is picked up from, and what happens to the -chain when the program is missing, slow or unhappy. The OpenRouter path is the -one that was always there and is checked here only for still being taken. +The CLIs are faked at subprocess.Popen: what the tests read is the argument +list each one is given, where the answer is picked up from, and what happens to +the chain when the program is missing, slow or unhappy. The OpenRouter path is +the one that was always there and is checked here only for still being taken. """ import os @@ -18,18 +18,26 @@ from tests.support import DikteTest, fake_urlopen, sent_json, url_error from tests.test_api import FakeServer, chat_reply -def fake_run(stdout="", code=0, stderr="", last_message=""): - """Stand in for subprocess.run, writing the file Codex would have written.""" +def fake_cli(stdout="", code=0, stderr="", last_message=""): + """Stand in for subprocess.Popen. + + _output hands the process a temporary file for each stream, so the fake + writes into those, plus the file Codex would have written on its way out. + """ calls = [] - def run(cmd, **kwargs): + def popen(cmd, **kwargs): calls.append(cmd) + kwargs["stdout"].write(stdout.encode("utf-8")) + kwargs["stderr"].write(stderr.encode("utf-8")) if last_message and "-o" in cmd: with open(cmd[cmd.index("-o") + 1], "w", encoding="utf-8") as fh: fh.write(last_message) - return subprocess.CompletedProcess(cmd, code, stdout, stderr) + proc = mock.Mock() + proc.returncode = code + return proc - return mock.patch.object(subprocess, "run", side_effect=run), calls + return mock.patch.object(subprocess, "Popen", side_effect=popen), calls class Provider(DikteTest): @@ -80,7 +88,7 @@ class OpenRouter(DikteTest): def test_no_cli_is_started_for_it(self): conf = self.config(openrouter_api_key="sk-or-test") - patcher, calls = fake_run(stdout="never") + patcher, calls = fake_cli(stdout="never") with patcher, mock.patch.object(api, "cleanup", return_value="Done."): cleanup.run("uh, done", conf, "the rules") self.assertEqual(calls, []) @@ -93,7 +101,7 @@ class ClaudeCode(DikteTest): self.patch_attr(cleanup.shutil, "which", lambda name: f"/usr/bin/{name}") def run_cleanup(self, text="uh, book it", **kwargs): - patcher, calls = fake_run(**kwargs) + patcher, calls = fake_cli(**kwargs) with patcher: answer = cleanup.run(text, self.conf, "the rules") return answer, calls[0] @@ -146,14 +154,18 @@ class ClaudeCode(DikteTest): self.run_cleanup(stdout="Book it.") self.assertIn("claude", str(caught.exception)) - def test_a_run_that_never_ends(self): - def run(cmd, **kwargs): - raise subprocess.TimeoutExpired(cmd, 180) + def test_a_run_that_never_ends_is_killed_with_its_whole_tree(self): + def popen(cmd, **kwargs): + proc = mock.Mock() + proc.wait.side_effect = subprocess.TimeoutExpired(cmd, 180) + return proc - with mock.patch.object(subprocess, "run", side_effect=run): + with mock.patch.object(subprocess, "Popen", side_effect=popen), \ + mock.patch.object(cleanup.assistant, "kill_tree") as kill: with self.assertRaises(cleanup.CleanupError) as caught: cleanup.run("uh, book it", self.conf, "the rules") self.assertIn("180", str(caught.exception)) + kill.assert_called_once() class Codex(DikteTest): @@ -163,7 +175,7 @@ class Codex(DikteTest): self.patch_attr(cleanup.shutil, "which", lambda name: f"/usr/bin/{name}") def run_cleanup(self, text="uh, book it", **kwargs): - patcher, calls = fake_run(**kwargs) + patcher, calls = fake_cli(**kwargs) with patcher: answer = cleanup.run(text, self.conf, "the rules") return answer, calls[0] @@ -280,7 +292,8 @@ class Here(DikteTest): self.assertIn("out of memory", str(caught.exception)) def test_no_cli_is_started_for_it(self): - patcher, calls = fake_run(stdout="never") + patcher, calls = fake_cli(stdout="never") with patcher, fake_urlopen(chat_reply("Done.")): cleanup.run("uh, done", self.conf, "the rules") self.assertEqual(calls, []) +