diff --git a/agent_runtime/execution_journal.py b/agent_runtime/execution_journal.py new file mode 100644 index 0000000..500b72a --- /dev/null +++ b/agent_runtime/execution_journal.py @@ -0,0 +1,298 @@ +"""Crash-safe execution journal for provider-requested tool calls.""" + +from __future__ import annotations + +import hashlib +import json +import uuid +from typing import Dict, List, Optional, Sequence + +from moonshine.utils import read_jsonl, shorten, utc_now + + +TOOL_EXECUTION_STARTED = "tool_execution_started" +TOOL_EXECUTION_FINISHED = "tool_execution_finished" +TOOL_EXECUTION_AMBIGUOUS = "tool_execution_ambiguous" +TOOL_EXECUTION_BLOCKED = "tool_execution_blocked" + + +class ToolExecutionJournal(object): + """Persist tool dispatch boundaries and fail closed after interrupted turns. + + The journal does not claim exactly-once execution. Instead it establishes a + conservative contract: write intent before dispatch, write a terminal marker + before the next call begins, and never automatically dispatch more tools in a + session when an earlier tool execution or tool-bearing turn has ambiguous + completion. + """ + + def __init__(self, runtime: Dict[str, object]): + self.runtime = runtime + self.store = runtime.get("session_store") if isinstance(runtime, dict) else None + self.session_id = str(runtime.get("session_id") or "").strip() if isinstance(runtime, dict) else "" + self.paths = getattr(self.store, "paths", None) if self.store is not None else None + self.turn_sequence, self.open_turn_sequences = self._turn_lifecycle() + + @property + def enabled(self) -> bool: + """Return whether durable session journaling is available.""" + return bool(self.store is not None and self.session_id and hasattr(self.store, "append_conversation_event")) + + def _render_json(self, value: object) -> str: + """Render a deterministic JSON-ish representation for hashes and previews.""" + return json.dumps(value, ensure_ascii=False, sort_keys=True, default=str) + + def _fingerprint(self, value: object, preview_chars: int = 800) -> Dict[str, str]: + """Return bounded trace metadata without duplicating large tool payloads.""" + rendered = self._render_json(value) + return { + "sha256": hashlib.sha256(rendered.encode("utf-8")).hexdigest(), + "preview": shorten(rendered, preview_chars), + } + + def _turn_lifecycle(self): + """Return monotonically numbered turn starts and currently open sequences.""" + if self.paths is None or not self.session_id: + return 0, [] + turn_events = [ + item + for item in read_jsonl(self.paths.session_turn_events_file(self.session_id)) + if isinstance(item, dict) and str(item.get("type") or "") in {"turn_started", "turn_completed"} + ] + sequence = 0 + open_sequences: List[int] = [] + for item in turn_events: + if str(item.get("type") or "") == "turn_started": + sequence += 1 + open_sequences.append(sequence) + elif open_sequences: + # A later resumed turn can complete while an older interrupted + # turn remains unresolved, so pair completion with the newest start. + open_sequences.pop() + return sequence, open_sequences + + def _payload(self, event: Dict[str, object]) -> Dict[str, object]: + payload = event.get("payload") or {} + return dict(payload) if isinstance(payload, dict) else {} + + def _events(self) -> List[Dict[str, object]]: + if self.store is None or not self.session_id or not hasattr(self.store, "get_conversation_events"): + return [] + return list(self.store.get_conversation_events(self.session_id)) + + def blockers(self) -> List[Dict[str, object]]: + """Return unresolved executions plus any prior interrupted tool-bearing turn.""" + events = self._events() + active: Dict[str, Dict[str, object]] = {} + for event in events: + kind = str(event.get("event_kind") or "") + if kind not in {TOOL_EXECUTION_STARTED, TOOL_EXECUTION_FINISHED, TOOL_EXECUTION_AMBIGUOUS}: + continue + payload = self._payload(event) + execution_id = str(payload.get("execution_id") or "").strip() + if not execution_id: + continue + if kind == TOOL_EXECUTION_FINISHED: + active.pop(execution_id, None) + continue + record = dict(payload) + record["state"] = "ambiguous" if kind == TOOL_EXECUTION_AMBIGUOUS else "started" + record["event_id"] = event.get("id") + active[execution_id] = record + + blockers = list(active.values()) + prior_turn = self._prior_interrupted_tool_turn(events) + if prior_turn is not None: + blockers.append(prior_turn) + return blockers + + def _prior_interrupted_tool_turn(self, events: Sequence[Dict[str, object]]) -> Optional[Dict[str, object]]: + """Return a prior open turn that executed tools before a later turn began.""" + # During a normal dispatch the current turn itself is open. Older open + # sequences represent turns that survived into a later user turn. + if len(self.open_turn_sequences) <= 1: + return None + prior_sequences = set(self.open_turn_sequences[:-1]) + for event in events: + if str(event.get("event_kind") or "") != TOOL_EXECUTION_STARTED: + continue + payload = self._payload(event) + try: + turn_sequence = int(payload.get("turn_sequence") or 0) + except (TypeError, ValueError): + turn_sequence = 0 + if turn_sequence not in prior_sequences: + continue + return { + "state": "interrupted_turn", + "execution_id": str(payload.get("execution_id") or "turn:%s" % turn_sequence), + "tool": str(payload.get("tool") or "unknown"), + "call_id": str(payload.get("call_id") or ""), + "turn_sequence": turn_sequence, + } + return None + + def _append(self, event_kind: str, content: str, payload: Dict[str, object]) -> None: + if not self.enabled: + return + self.store.append_conversation_event( + self.session_id, + event_kind=event_kind, + role="tool", + content=content, + payload=dict(payload), + ) + + def _mark_session_interrupted(self, interruption: Dict[str, object]) -> None: + if self.store is None or not self.session_id: + return + now = utc_now() + if hasattr(self.store, "update_session_meta"): + self.store.update_session_meta( + self.session_id, + status="interrupted", + updated_at=now, + interrupted_tool_execution=dict(interruption), + ) + db = getattr(self.store, "db", None) + if db is not None and hasattr(db, "update_session"): + db.update_session(self.session_id, updated_at=now, status="interrupted") + + def begin(self, call: object) -> str: + """Write a durable intent immediately before dispatch.""" + execution_id = "tool-exec-%s" % uuid.uuid4().hex[:12] + tool_name = str(getattr(call, "name", "") or "") + call_id = str(getattr(call, "call_id", "") or "") + arguments = dict(getattr(call, "arguments", {}) or {}) + arguments_fingerprint = self._fingerprint(arguments) + payload = { + "execution_id": execution_id, + "tool": tool_name, + "call_id": call_id, + "arguments_sha256": arguments_fingerprint["sha256"], + "arguments_preview": arguments_fingerprint["preview"], + "turn_sequence": self.turn_sequence, + "tool_round": self.runtime.get("_current_tool_round", ""), + "started_at": utc_now(), + } + self._append( + TOOL_EXECUTION_STARTED, + "Tool execution started: %s (%s)" % (tool_name, call_id or execution_id), + payload, + ) + return execution_id + + def finish(self, call: object, execution_id: str, *, output: object, error: Optional[str]) -> None: + """Write a terminal marker before the next call is dispatched.""" + output_fingerprint = self._fingerprint(output, preview_chars=1200) + payload = { + "execution_id": execution_id, + "tool": str(getattr(call, "name", "") or ""), + "call_id": str(getattr(call, "call_id", "") or ""), + "turn_sequence": self.turn_sequence, + "tool_round": self.runtime.get("_current_tool_round", ""), + "outcome": "error" if error else "ok", + "error": shorten(str(error or ""), 500), + "output_preview": output_fingerprint["preview"], + "output_sha256": output_fingerprint["sha256"], + "finished_at": utc_now(), + } + self._append( + TOOL_EXECUTION_FINISHED, + "Tool execution finished: %s (%s)" % (payload["tool"], payload["call_id"] or execution_id), + payload, + ) + + def mark_ambiguous(self, call: object, execution_id: str, exc: BaseException) -> None: + """Record a process-level interruption whose completion is unknowable.""" + arguments_fingerprint = self._fingerprint(dict(getattr(call, "arguments", {}) or {})) + payload = { + "execution_id": execution_id, + "tool": str(getattr(call, "name", "") or ""), + "call_id": str(getattr(call, "call_id", "") or ""), + "arguments_sha256": arguments_fingerprint["sha256"], + "arguments_preview": arguments_fingerprint["preview"], + "turn_sequence": self.turn_sequence, + "tool_round": self.runtime.get("_current_tool_round", ""), + "state": "ambiguous", + "interruption_type": type(exc).__name__, + "interruption": shorten(str(exc), 500), + "interrupted_at": utc_now(), + } + self._append( + TOOL_EXECUTION_AMBIGUOUS, + "Tool execution became ambiguous after interruption: %s (%s)" + % (payload["tool"], payload["call_id"] or execution_id), + payload, + ) + self._mark_session_interrupted( + { + "execution_id": execution_id, + "tool": payload["tool"], + "call_id": payload["call_id"], + "state": "ambiguous", + } + ) + + def blocked_results(self, calls: List[object], blockers: Sequence[Dict[str, object]]) -> List[Dict[str, object]]: + """Return provider-visible errors without dispatching any requested tool.""" + first = dict(blockers[0]) if blockers else {} + blocker_tool = str(first.get("tool") or "unknown") + blocker_call_id = str(first.get("call_id") or "unknown") + blocker_execution_id = str(first.get("execution_id") or "unknown") + blocker_state = str(first.get("state") or "ambiguous") + self._mark_session_interrupted( + { + "execution_id": blocker_execution_id, + "tool": blocker_tool, + "call_id": blocker_call_id, + "state": blocker_state, + } + ) + reason = ( + "a prior tool-bearing turn was interrupted before Moonshine durably completed the turn" + if blocker_state == "interrupted_turn" + else "a prior tool execution has ambiguous completion" + ) + message = ( + "Tool dispatch is blocked because %s: tool=%s, call_id=%s, execution_id=%s. " + "Moonshine will not replay or dispatch additional tools automatically because prior " + "handlers may already have produced external side effects. Inspect the session records " + "and continue in a fresh session once the ambiguity is resolved." + % (reason, blocker_tool, blocker_call_id, blocker_execution_id) + ) + public_blockers = [ + { + key: item.get(key) + for key in ("state", "execution_id", "tool", "call_id", "turn_sequence", "started_at", "interrupted_at") + if item.get(key) not in {None, ""} + } + for item in blockers + ] + results: List[Dict[str, object]] = [] + for call in calls: + result = { + "name": getattr(call, "name", ""), + "call_id": getattr(call, "call_id", ""), + "arguments": getattr(call, "arguments", {}), + "output": { + "status": "blocked_interrupted_execution", + "message": message, + "ambiguous_executions": public_blockers, + }, + "error": message, + } + results.append(result) + self.runtime.setdefault("_tool_results_in_round", []).append(result) + self._append( + TOOL_EXECUTION_BLOCKED, + "Blocked tool dispatch: %s" % str(getattr(call, "name", "") or ""), + { + "tool": str(getattr(call, "name", "") or ""), + "call_id": str(getattr(call, "call_id", "") or ""), + "blocked_at": utc_now(), + "blocker_states": [str(item.get("state") or "") for item in blockers], + "ambiguous_execution_ids": [str(item.get("execution_id") or "") for item in blockers], + }, + ) + return results diff --git a/model_tools.py b/model_tools.py index 8674ecc..217efaa 100644 --- a/model_tools.py +++ b/model_tools.py @@ -5,6 +5,9 @@ import traceback from typing import Dict, List, Optional, Sequence +from moonshine.agent_runtime.execution_journal import ToolExecutionJournal + + def collect_tool_schemas( registry, mode: Optional[str] = None, @@ -17,26 +20,43 @@ def collect_tool_schemas( def handle_function_calls(registry, calls: List[object], runtime: Dict[str, object]) -> List[Dict[str, object]]: - """Dispatch provider tool calls through the registry.""" + """Dispatch provider tool calls through the registry with crash-safe journaling.""" + journal = ToolExecutionJournal(runtime) + blockers = journal.blockers() + if blockers: + return journal.blocked_results(calls, blockers) + results = [] for call in calls: + execution_id = journal.begin(call) try: - result = registry.dispatch(call.name, call.arguments, runtime) - error = None - except Exception as exc: - result = { - "error": str(exc), - "traceback": traceback.format_exc(limit=3), - } - error = str(exc) - results.append( - { - "name": call.name, - "call_id": getattr(call, "call_id", ""), - "arguments": call.arguments, - "output": result, - "error": error, - } - ) - runtime.setdefault("_tool_results_in_round", []).append(results[-1]) + try: + result = registry.dispatch(call.name, call.arguments, runtime) + error = None + except Exception as exc: + result = { + "error": str(exc), + "traceback": traceback.format_exc(limit=3), + } + error = str(exc) + journal.finish(call, execution_id, output=result, error=error) + except BaseException as exc: + try: + journal.mark_ambiguous(call, execution_id, exc) + except Exception: + # Never replace the process-level interruption with a best-effort + # journaling failure. A durable start record, when it was written, + # is itself enough for the next process to fail closed. + pass + raise + + result_record = { + "name": call.name, + "call_id": getattr(call, "call_id", ""), + "arguments": call.arguments, + "output": result, + "error": error, + } + results.append(result_record) + runtime.setdefault("_tool_results_in_round", []).append(result_record) return results diff --git a/tests/test_interrupted_tool_recovery.py b/tests/test_interrupted_tool_recovery.py new file mode 100644 index 0000000..2acab43 --- /dev/null +++ b/tests/test_interrupted_tool_recovery.py @@ -0,0 +1,305 @@ +"""Regression tests for crash-safe tool execution journaling.""" + +from __future__ import annotations + +import tempfile +import unittest +from pathlib import Path + +from moonshine.agent_runtime.execution_journal import ( + TOOL_EXECUTION_AMBIGUOUS, + TOOL_EXECUTION_BLOCKED, + TOOL_EXECUTION_FINISHED, + TOOL_EXECUTION_STARTED, +) +from moonshine.model_tools import handle_function_calls +from moonshine.moonshine_constants import MoonshinePaths +from moonshine.providers import ProviderToolCall +from moonshine.storage.session_store import SessionStore + + +class ScriptedRegistry(object): + """Minimal deterministic registry for execution-lifecycle tests.""" + + def __init__(self, handlers): + self.handlers = dict(handlers) + self.dispatches = [] + + def dispatch(self, name, arguments, runtime): + self.dispatches.append((name, dict(arguments or {}))) + handler = self.handlers[name] + return handler(runtime, **dict(arguments or {})) + + +class InterruptedToolRecoveryTest(unittest.TestCase): + def setUp(self): + self.temp_dir = tempfile.TemporaryDirectory() + self.addCleanup(self.temp_dir.cleanup) + self.paths = MoonshinePaths(Path(self.temp_dir.name)) + self.store = SessionStore(self.paths) + self.session_id = self.store.create_session("chat", "tool-recovery-test") + + def _runtime(self, store=None): + return { + "session_store": store or self.store, + "session_id": self.session_id, + "_tool_results_in_round": [], + } + + def _execution_events(self, store=None): + store = store or self.store + return [ + item + for item in store.get_conversation_events(self.session_id) + if str(item.get("event_kind") or "").startswith("tool_execution_") + ] + + def test_completed_call_is_terminal_before_later_call_is_interrupted(self): + side_effects = [] + + def complete(runtime, value): + side_effects.append("complete:%s" % value) + return {"value": value, "status": "done"} + + def interrupt(runtime): + side_effects.append("interrupt-side-effect") + raise KeyboardInterrupt("simulated process interruption") + + registry = ScriptedRegistry({"complete": complete, "interrupt": interrupt}) + runtime = self._runtime() + calls = [ + ProviderToolCall(name="complete", arguments={"value": 7}, call_id="call-complete"), + ProviderToolCall(name="interrupt", arguments={}, call_id="call-interrupt"), + ] + + with self.assertRaises(KeyboardInterrupt): + handle_function_calls(registry, calls, runtime) + + self.assertEqual(side_effects, ["complete:7", "interrupt-side-effect"]) + self.assertEqual([item["name"] for item in runtime["_tool_results_in_round"]], ["complete"]) + + events = self._execution_events() + complete_started = [ + item for item in events + if item["event_kind"] == TOOL_EXECUTION_STARTED + and item["payload"].get("call_id") == "call-complete" + ] + complete_finished = [ + item for item in events + if item["event_kind"] == TOOL_EXECUTION_FINISHED + and item["payload"].get("call_id") == "call-complete" + ] + interrupted_started = [ + item for item in events + if item["event_kind"] == TOOL_EXECUTION_STARTED + and item["payload"].get("call_id") == "call-interrupt" + ] + interrupted_ambiguous = [ + item for item in events + if item["event_kind"] == TOOL_EXECUTION_AMBIGUOUS + and item["payload"].get("call_id") == "call-interrupt" + ] + + self.assertEqual(len(complete_started), 1) + self.assertEqual(len(complete_finished), 1) + self.assertEqual( + complete_started[0]["payload"]["execution_id"], + complete_finished[0]["payload"]["execution_id"], + ) + self.assertEqual(complete_finished[0]["payload"]["outcome"], "ok") + self.assertTrue(complete_finished[0]["payload"]["output_sha256"]) + self.assertIn('"status": "done"', complete_finished[0]["payload"]["output_preview"]) + self.assertNotIn("arguments", complete_started[0]["payload"]) + self.assertIn("arguments_sha256", complete_started[0]["payload"]) + + self.assertEqual(len(interrupted_started), 1) + self.assertEqual(len(interrupted_ambiguous), 1) + self.assertEqual( + interrupted_started[0]["payload"]["execution_id"], + interrupted_ambiguous[0]["payload"]["execution_id"], + ) + self.assertEqual(interrupted_ambiguous[0]["payload"]["state"], "ambiguous") + self.assertEqual(self.store.get_session_meta(self.session_id)["status"], "interrupted") + + def test_restart_blocks_new_tool_dispatch_after_orphaned_execution_start(self): + execution_id = "tool-exec-hard-crash" + self.store.append_conversation_event( + self.session_id, + event_kind=TOOL_EXECUTION_STARTED, + role="tool", + content="Tool execution started before simulated hard crash", + payload={ + "execution_id": execution_id, + "tool": "external_side_effect", + "call_id": "call-before-crash", + "arguments_sha256": "deadbeef", + "arguments_preview": '{"value": 1}', + "turn_sequence": 0, + }, + ) + + restarted_store = SessionStore(self.paths) + dispatch_count = [] + + def must_not_run(runtime): + dispatch_count.append(1) + return {"unexpected": True} + + registry = ScriptedRegistry({"must_not_run": must_not_run}) + runtime = self._runtime(restarted_store) + results = handle_function_calls( + registry, + [ProviderToolCall(name="must_not_run", arguments={}, call_id="call-after-restart")], + runtime, + ) + + self.assertEqual(dispatch_count, []) + self.assertEqual(registry.dispatches, []) + self.assertEqual(len(results), 1) + self.assertEqual(results[0]["output"]["status"], "blocked_interrupted_execution") + self.assertIn(execution_id, results[0]["error"]) + self.assertEqual(restarted_store.get_session_meta(self.session_id)["status"], "interrupted") + blocked = [ + item for item in self._execution_events(restarted_store) + if item["event_kind"] == TOOL_EXECUTION_BLOCKED + ] + self.assertEqual(len(blocked), 1) + self.assertIn(execution_id, blocked[0]["payload"]["ambiguous_execution_ids"]) + + def test_completed_resume_turn_does_not_hide_older_interrupted_tool_turn(self): + # Turn 1 starts and executes a tool, but the process dies before its + # turn_completed record is written. + self.store.append_turn_event( + self.session_id, + {"type": "turn_started", "text": "first", "created_at": "2026-09-05T00:00:00Z"}, + ) + execution_id = "tool-exec-finished-before-crash" + self.store.append_conversation_event( + self.session_id, + event_kind=TOOL_EXECUTION_STARTED, + role="tool", + content="Tool execution started", + payload={ + "execution_id": execution_id, + "tool": "external_side_effect", + "call_id": "call-first-turn", + "turn_sequence": 1, + "arguments_sha256": "abc", + "arguments_preview": "{}", + }, + ) + self.store.append_conversation_event( + self.session_id, + event_kind=TOOL_EXECUTION_FINISHED, + role="tool", + content="Tool execution finished", + payload={ + "execution_id": execution_id, + "tool": "external_side_effect", + "call_id": "call-first-turn", + "turn_sequence": 1, + "outcome": "ok", + "output_sha256": "def", + "output_preview": '{"ok": true}', + }, + ) + + # Turn 2 is a resumed, tool-free turn that completes. LIFO pairing must + # close turn 2, not accidentally consume the older interrupted turn 1. + self.store.append_turn_event( + self.session_id, + {"type": "turn_started", "text": "resume-one", "created_at": "2026-09-05T00:00:00Z"}, + ) + self.store.append_turn_event( + self.session_id, + {"type": "turn_completed", "text": "resume-one done", "created_at": "2026-09-05T00:00:00Z"}, + ) + # Turn 3 begins and attempts another tool call. + self.store.append_turn_event( + self.session_id, + {"type": "turn_started", "text": "resume-two", "created_at": "2026-09-05T00:00:00Z"}, + ) + + dispatch_count = [] + + def must_not_run(runtime): + dispatch_count.append(1) + return {"unexpected": True} + + registry = ScriptedRegistry({"must_not_run": must_not_run}) + results = handle_function_calls( + registry, + [ProviderToolCall(name="must_not_run", arguments={}, call_id="call-resumed")], + self._runtime(), + ) + + self.assertEqual(dispatch_count, []) + self.assertEqual(registry.dispatches, []) + self.assertEqual(results[0]["output"]["status"], "blocked_interrupted_execution") + self.assertIn("prior tool-bearing turn was interrupted", results[0]["error"]) + blocker = results[0]["output"]["ambiguous_executions"][0] + self.assertEqual(blocker["state"], "interrupted_turn") + self.assertEqual(blocker["turn_sequence"], 1) + + def test_ordinary_tool_error_is_terminal_and_does_not_poison_future_dispatch(self): + def fail(runtime): + raise RuntimeError("deterministic tool failure") + + failing_registry = ScriptedRegistry({"fail": fail}) + first = handle_function_calls( + failing_registry, + [ProviderToolCall(name="fail", arguments={}, call_id="call-fail")], + self._runtime(), + ) + self.assertEqual(len(first), 1) + self.assertIn("deterministic tool failure", first[0]["error"]) + + finish_events = [ + item for item in self._execution_events() + if item["event_kind"] == TOOL_EXECUTION_FINISHED + and item["payload"].get("call_id") == "call-fail" + ] + self.assertEqual(len(finish_events), 1) + self.assertEqual(finish_events[0]["payload"]["outcome"], "error") + + later_effects = [] + + def succeed(runtime): + later_effects.append("ran") + return {"ok": True} + + succeeding_registry = ScriptedRegistry({"succeed": succeed}) + second = handle_function_calls( + succeeding_registry, + [ProviderToolCall(name="succeed", arguments={}, call_id="call-succeed")], + self._runtime(), + ) + self.assertEqual(later_effects, ["ran"]) + self.assertIsNone(second[0]["error"]) + + def test_dispatch_without_session_store_keeps_legacy_result_shape(self): + effects = [] + + def succeed(runtime, value): + effects.append(value) + return {"value": value} + + registry = ScriptedRegistry({"succeed": succeed}) + runtime = {"_tool_results_in_round": []} + results = handle_function_calls( + registry, + [ProviderToolCall(name="succeed", arguments={"value": 3}, call_id="call-no-store")], + runtime, + ) + + self.assertEqual(effects, [3]) + self.assertEqual(results[0]["output"], {"value": 3}) + self.assertIsNone(results[0]["error"]) + self.assertEqual( + set(results[0]), + {"name", "call_id", "arguments", "output", "error"}, + ) + + +if __name__ == "__main__": + unittest.main()