fix(session): persist tool activity before projection

This commit is contained in:
elcocoel 2026-07-27 04:09:41 +01:00 committed by kshitij
parent e643f2e910
commit 858bedea02
4 changed files with 385 additions and 71 deletions

View file

@ -1061,6 +1061,9 @@ def run_conversation(
# Commentary deduplication spans all provider continuations and tool calls
# within one user turn, but must not suppress the same phrase next turn.
agent._delivered_interim_texts = set()
# A configured SessionDB append failure halts only the affected turn. A
# cached gateway agent must recover on the next message if storage did.
agent._incremental_persistence_failed = False
# Main conversation loop counters (pure locals consumed by the loop below).
api_call_count = 0
@ -5785,8 +5788,6 @@ def run_conversation(
and previous_interim_visible == current_interim_visible
)
messages.append(assistant_msg)
if not duplicate_previous_interim:
agent._emit_interim_assistant_message(assistant_msg)
# Mixed batch: error-result the invalid calls and strip them
# from the execution set. The assistant message above keeps
@ -5808,13 +5809,17 @@ def run_conversation(
if tc.function.name in agent.valid_tool_names
]
_tool_turn_persisted = None
try:
# Persist the assistant tool-call turn before any tool
# side effects run. If a destructive tool restarts or
# terminates Hermes mid-turn, resume logic still sees the
# exact tool-call block that already executed.
agent._flush_messages_to_session_db(messages, conversation_history)
_tool_turn_persisted = agent._flush_messages_to_session_db(
messages, conversation_history
)
except Exception as exc:
_tool_turn_persisted = False
logger.warning(
"Incremental tool-call persistence failed before execution "
"(session=%s): %s",
@ -5822,6 +5827,22 @@ def run_conversation(
exc,
)
if _tool_turn_persisted is False:
# The canonical append failed. Do not project the row or
# run side-effecting tools from state that exists only in
# this process. Breaking also avoids retrying the same
# unpersisted turn until the iteration budget is exhausted.
_turn_exit_reason = "session_persistence_failed"
final_response = ""
failed = True
break
# A UI must never observe an assistant/tool-call row that is
# still only an ephemeral in-memory projection. Emit interim
# commentary only after the canonical SessionDB append above.
if not duplicate_previous_interim:
agent._emit_interim_assistant_message(assistant_msg)
# Close any open streaming display (response box, reasoning
# box) before tool execution begins. Intermediate turns may
# have streamed early content that opened the response box;
@ -5836,6 +5857,15 @@ def run_conversation(
agent._execute_tool_calls(assistant_message, messages, effective_task_id, api_call_count)
if getattr(agent, "_incremental_persistence_failed", False):
# A tool result could not be made canonical. Do not send
# the in-memory result back to the model or project any
# later events from this turn.
_turn_exit_reason = "session_persistence_failed"
final_response = ""
failed = True
break
if agent._tool_guardrail_halt_decision is not None:
decision = agent._tool_guardrail_halt_decision
_turn_exit_reason = "guardrail_halt"

View file

@ -140,8 +140,8 @@ def _flush_session_db_after_tool_progress(
messages: list,
*,
stage: str,
) -> None:
"""Best-effort incremental SessionDB flush for tool-call progress.
) -> bool:
"""Flush tool-call progress before projecting it to any UI surface.
Tool execution can perform side effects that terminate or restart the
current Hermes process before the normal turn-end persistence path runs.
@ -149,9 +149,14 @@ def _flush_session_db_after_tool_progress(
transcript survives destructive-but-valid tool calls.
"""
try:
agent._flush_messages_to_session_db(messages)
persisted = agent._flush_messages_to_session_db(messages) is not False
if not persisted:
agent._incremental_persistence_failed = True
return persisted
except Exception as exc:
agent._incremental_persistence_failed = True
logger.warning("Incremental tool-call persistence failed after %s: %s", stage, exc)
return False
def _ra():
@ -861,6 +866,8 @@ def execute_tool_calls_concurrent(agent, assistant_message, messages: list, effe
for i, (tc, name, args, middleware_trace, block_result, blocked_by_guardrail) in enumerate(parsed_calls):
r = results[i]
blocked = False
is_error = True
progress_function_name = name
# A worker can finish and write results[i] in the window between the
# deadline snapshot (timed_out_indices, taken from not_done) and this
# loop. Prefer that real result over a fabricated timeout message — the
@ -916,6 +923,7 @@ def execute_tool_calls_concurrent(agent, assistant_message, messages: list, effe
tool_duration = 0.0
else:
function_name, function_args, function_result, tool_duration, is_error, blocked, middleware_trace = r
progress_function_name = function_name
if blocked:
effect_disposition = "none"
@ -943,44 +951,15 @@ def execute_tool_calls_concurrent(agent, assistant_message, messages: list, effe
except Exception as _ver_err:
logging.debug("file-mutation verifier record failed: %s", _ver_err)
if not blocked and agent.tool_progress_callback:
try:
agent.tool_progress_callback(
"tool.completed", function_name, None, None,
duration=tool_duration, is_error=is_error,
result=function_result,
)
except Exception as cb_err:
logging.debug(f"Tool progress callback error: {cb_err}")
if agent.verbose_logging:
logging.debug(f"Tool {function_name} completed in {tool_duration:.2f}s")
logging.debug(f"Tool result ({len(function_result)} chars): {function_result}")
# Print cute message per tool
if agent._should_emit_quiet_tool_messages():
cute_msg = _get_cute_tool_message_impl(name, args, tool_duration, result=function_result)
agent._safe_print(f" {cute_msg}")
elif not agent.quiet_mode and getattr(agent, "tool_progress_mode", "all") != "off":
_preview_str = _multimodal_text_summary(function_result)
if agent.verbose_logging:
print(f" ✅ Tool {i+1} completed in {tool_duration:.2f}s")
print(agent._wrap_verbose("Result: ", _preview_str))
else:
response_preview = _preview_str[:agent.log_prefix_chars] + "..." if len(_preview_str) > agent.log_prefix_chars else _preview_str
print(f" ✅ Tool {i+1} completed in {tool_duration:.2f}s - {response_preview}")
agent._current_tool = None
_status_suffix = " (error)" if is_error else ""
agent._touch_activity(f"tool completed: {name} ({tool_duration:.1f}s){_status_suffix}")
if not blocked and agent.tool_complete_callback:
try:
display_args = _redact_tool_args_for_display(name, args) or args
agent.tool_complete_callback(tc.id, name, display_args, function_result)
except Exception as cb_err:
logging.debug(f"Tool complete callback error: {cb_err}")
display_function_result = function_result
function_result = maybe_persist_tool_result(
content=function_result,
tool_name=name,
@ -1015,6 +994,50 @@ def execute_tool_calls_concurrent(agent, assistant_message, messages: list, effe
)
messages.append(tool_message)
risk_metadata = tool_message.get("_tool_output_risk")
if not _flush_session_db_after_tool_progress(
agent,
messages,
stage=f"tool result {name}",
):
return
# Every completion surface is downstream of the canonical append. If
# the UI bridge or process dies while projecting one of these events,
# resume can reconstruct the tool result that was already visible.
if not blocked and agent.tool_progress_callback:
try:
agent.tool_progress_callback(
"tool.completed", progress_function_name, None, None,
duration=tool_duration, is_error=is_error,
result=display_function_result,
)
except Exception as cb_err:
logging.debug(f"Tool progress callback error: {cb_err}")
# Print cute message per tool
if agent._should_emit_quiet_tool_messages():
cute_msg = _get_cute_tool_message_impl(
name, args, tool_duration, result=display_function_result,
)
agent._safe_print(f" {cute_msg}")
elif not agent.quiet_mode and getattr(agent, "tool_progress_mode", "all") != "off":
_preview_str = _multimodal_text_summary(display_function_result)
if agent.verbose_logging:
print(f" ✅ Tool {i+1} completed in {tool_duration:.2f}s")
print(agent._wrap_verbose("Result: ", _preview_str))
else:
response_preview = _preview_str[:agent.log_prefix_chars] + "..." if len(_preview_str) > agent.log_prefix_chars else _preview_str
print(f" ✅ Tool {i+1} completed in {tool_duration:.2f}s - {response_preview}")
if not blocked and agent.tool_complete_callback:
try:
display_args = _redact_tool_args_for_display(name, args) or args
agent.tool_complete_callback(
tc.id, name, display_args, display_function_result,
)
except Exception as cb_err:
logging.debug(f"Tool complete callback error: {cb_err}")
if (
risk_metadata is not None
and risk_metadata.get("risk") != "low"
@ -1031,11 +1054,6 @@ def execute_tool_calls_concurrent(agent, assistant_message, messages: list, effe
)
except Exception as cb_err:
logging.debug("Tool output risk callback error: %s", cb_err)
_flush_session_db_after_tool_progress(
agent,
messages,
stage=f"tool result {name}",
)
# ── Per-tool /steer drain ───────────────────────────────────
# Same as the sequential path: drain between each collected
@ -1067,6 +1085,8 @@ def execute_tool_calls_sequential(agent, assistant_message, messages: list, effe
# Resolve the context-scaled tool-output budget once per turn.
_tool_budget = _budget_for_agent(agent)
for i, tool_call in enumerate(assistant_message.tool_calls, 1):
if getattr(agent, "_incremental_persistence_failed", False):
return
# SAFETY: check interrupt BEFORE starting each tool.
# If the user sent "stop" during a previous tool's execution,
# do NOT start any more tools -- skip them all immediately.
@ -1082,11 +1102,12 @@ def execute_tool_calls_sequential(agent, assistant_message, messages: list, effe
skipped_tc.id,
effect_disposition="none",
))
_flush_session_db_after_tool_progress(
if not _flush_session_db_after_tool_progress(
agent,
messages,
stage=f"cancelled tool result {skipped_name}",
)
):
return
break
function_name = tool_call.function.name
@ -1102,11 +1123,12 @@ def execute_tool_calls_sequential(agent, assistant_message, messages: list, effe
tool_call.id,
)
)
_flush_session_db_after_tool_progress(
if not _flush_session_db_after_tool_progress(
agent,
messages,
stage=f"invalid tool arguments {function_name}",
)
):
return
agent._apply_pending_steer_to_tool_results(messages, 1)
continue
@ -1670,16 +1692,6 @@ def execute_tool_calls_sequential(agent, assistant_message, messages: list, effe
except Exception as _ver_err:
logging.debug("file-mutation verifier record failed: %s", _ver_err)
if not _execution_blocked and agent.tool_progress_callback:
try:
agent.tool_progress_callback(
"tool.completed", function_name, None, None,
duration=tool_duration, is_error=_is_error_result,
result=function_result,
)
except Exception as cb_err:
logging.debug(f"Tool progress callback error: {cb_err}")
agent._current_tool = None
_status_suffix = " (error)" if _is_error_result else ""
agent._touch_activity(f"tool completed: {function_name} ({tool_duration:.1f}s){_status_suffix}")
@ -1689,13 +1701,7 @@ def execute_tool_calls_sequential(agent, assistant_message, messages: list, effe
_log_result = _multimodal_text_summary(function_result)
logging.debug(f"Tool result ({len(_log_result)} chars): {_log_result}")
if not _execution_blocked and agent.tool_complete_callback:
try:
display_args = _redact_tool_args_for_display(function_name, function_args) or function_args
agent.tool_complete_callback(tool_call.id, function_name, display_args, function_result)
except Exception as cb_err:
logging.debug(f"Tool complete callback error: {cb_err}")
display_function_result = function_result
function_result = maybe_persist_tool_result(
content=function_result,
tool_name=function_name,
@ -1718,6 +1724,40 @@ def execute_tool_calls_sequential(agent, assistant_message, messages: list, effe
tool_message = make_tool_result_message(function_name, _tool_content, tool_call.id)
messages.append(tool_message)
risk_metadata = tool_message.get("_tool_output_risk")
if not _flush_session_db_after_tool_progress(
agent,
messages,
stage=f"tool result {function_name}",
):
return
# UI completion/progress events are projections of the canonical tool
# row, never a competing in-memory authority.
if not _execution_blocked and agent.tool_progress_callback:
try:
agent.tool_progress_callback(
"tool.completed", function_name, None, None,
duration=tool_duration, is_error=_is_error_result,
result=display_function_result,
)
except Exception as cb_err:
logging.debug(f"Tool progress callback error: {cb_err}")
if not _execution_blocked and agent.tool_complete_callback:
try:
display_args = (
_redact_tool_args_for_display(function_name, function_args)
or function_args
)
agent.tool_complete_callback(
tool_call.id,
function_name,
display_args,
display_function_result,
)
except Exception as cb_err:
logging.debug(f"Tool complete callback error: {cb_err}")
if (
risk_metadata is not None
and risk_metadata.get("risk") != "low"
@ -1734,11 +1774,6 @@ def execute_tool_calls_sequential(agent, assistant_message, messages: list, effe
)
except Exception as cb_err:
logging.debug("Tool output risk callback error: %s", cb_err)
_flush_session_db_after_tool_progress(
agent,
messages,
stage=f"tool result {function_name}",
)
# ── Per-tool /steer drain ───────────────────────────────────
# Drain pending steer BETWEEN individual tool calls so the
@ -1766,11 +1801,12 @@ def execute_tool_calls_sequential(agent, assistant_message, messages: list, effe
skipped_tc.id,
effect_disposition="none",
))
_flush_session_db_after_tool_progress(
if not _flush_session_db_after_tool_progress(
agent,
messages,
stage=f"skipped tool result {skipped_name}",
)
):
return
break
if agent.tool_delay > 0 and i < len(assistant_message.tool_calls):
@ -1821,6 +1857,8 @@ def execute_tool_calls_segmented(agent, assistant_message, messages: list, effec
segments = _plan_tool_batch_segments(assistant_message.tool_calls, execution_cwd=_exec_cwd)
for kind, calls in segments:
if getattr(agent, "_incremental_persistence_failed", False):
return
segment_message = SimpleNamespace(tool_calls=list(calls))
if kind == "parallel":
execute_tool_calls_concurrent(
@ -1833,6 +1871,9 @@ def execute_tool_calls_segmented(agent, assistant_message, messages: list, effec
finalize=False,
)
if getattr(agent, "_incremental_persistence_failed", False):
return
# ── Whole-turn finalize (budget + /steer) ─────────────────────────
total_tools = len(assistant_message.tool_calls)
if total_tools > 0:

View file

@ -1918,9 +1918,9 @@ class AIAgent:
# where the next live turn re-reads it as an instruction and the agent
# "becomes" the curator. Hard-stop before any DB touch.
if getattr(self, "_persist_disabled", False):
return
return None
if not self._session_db:
return
return None
# Persist user-message override (#48677 chokepoint): historically this
# mutated the live `messages` list in place, which — on the early
# crash-resilience persist that runs BEFORE the API call is built —
@ -2122,8 +2122,10 @@ class AIAgent:
# allocated next turn at a recycled address.
self._flushed_db_message_ids = set()
self._last_flushed_db_idx = len(messages)
return True
except Exception as e:
logger.warning("Session DB append_message failed: %s", e)
return False
def _get_messages_up_to_last_assistant(self, messages: List[Dict]) -> List[Dict]:
"""

View file

@ -28,7 +28,12 @@ from pathlib import Path
import tempfile
from unittest.mock import MagicMock, patch
import pytest
from agent.tool_dispatch_helpers import make_tool_result_message
from agent.agent_runtime_helpers import sanitize_api_messages
from agent.tool_executor import execute_tool_calls_segmented
from hermes_state import SessionDB
from run_agent import AIAgent
@ -75,6 +80,31 @@ def _make_agent():
return agent
def _attach_real_session_db(agent, db_path: Path, session_id: str) -> SessionDB:
db = SessionDB(db_path=db_path)
db.create_session(session_id=session_id, source="tui", model="test/model")
agent._session_db = db
agent._session_db_created = True
agent.session_id = session_id
agent._last_flushed_db_idx = 0
agent._flushed_db_message_ids = set()
agent._flushed_db_message_session_id = None
agent._persist_disabled = False
return db
def _durable_messages(db_path: Path, session_id: str) -> list[dict]:
restarted_db = SessionDB(db_path=db_path)
try:
return restarted_db.get_messages_as_conversation(session_id)
finally:
restarted_db.close()
def _durable_roles(db_path: Path, session_id: str) -> list[str]:
return [message["role"] for message in _durable_messages(db_path, session_id)]
def _mock_tool_call(name="web_search", arguments="{}", call_id="call_1"):
return SimpleNamespace(
id=call_id,
@ -142,6 +172,78 @@ def test_run_conversation_flushes_assistant_tool_call_before_execution():
assert result["final_response"] == "done"
def test_interim_assistant_is_durable_before_ui_projection_on_abnormal_exit(tmp_path):
"""A visible interim assistant row must survive an immediate process exit.
``GeneratorExit`` models an uncatchable turn interruption at the UI bridge:
no turn finalizer or graceful shutdown persistence is allowed to rescue the
row after the callback observes it.
"""
agent = _make_agent()
db_path = tmp_path / "state.db"
session_id = "interim-abnormal-exit"
db = _attach_real_session_db(agent, db_path, session_id)
tool_call = _mock_tool_call(call_id="visible-call")
agent.client.chat.completions.create.return_value = _mock_response(
content="I'll inspect the repository now.",
finish_reason="tool_calls",
tool_calls=[tool_call],
)
roles_seen_by_ui: list[str] = []
def _ui_projection(_text, *, already_streamed=False):
roles_seen_by_ui.extend(_durable_roles(db_path, session_id))
raise GeneratorExit("simulated process termination after UI projection")
agent.interim_assistant_callback = _ui_projection
try:
with pytest.raises(GeneratorExit, match="simulated process termination"):
agent.run_conversation("inspect the repository")
finally:
db.close()
assert roles_seen_by_ui == ["user", "assistant"]
durable = _durable_messages(db_path, session_id)
assert [message["role"] for message in durable] == ["user", "assistant"]
assert durable[1]["content"] == "I'll inspect the repository now."
assert durable[1]["tool_calls"][0]["id"] == "visible-call"
# Cold-resume reconciliation closes the interrupted call in the provider
# payload without mutating or duplicating the canonical transcript.
resumed = sanitize_api_messages(durable)
assert [message["role"] for message in resumed] == [
"user",
"assistant",
"tool",
]
assert resumed[2]["tool_call_id"] == "visible-call"
assert len(_durable_messages(db_path, session_id)) == 2
def test_failed_assistant_persist_blocks_ui_projection_and_tool_side_effects():
agent = _make_agent()
tool_call = _mock_tool_call(call_id="must-not-run")
agent.client.chat.completions.create.return_value = _mock_response(
content="I'll inspect the repository now.",
finish_reason="tool_calls",
tool_calls=[tool_call],
)
agent._flush_messages_to_session_db = MagicMock(return_value=False)
agent.interim_assistant_callback = MagicMock()
agent._execute_tool_calls = MagicMock()
result = agent.run_conversation("inspect the repository")
agent.interim_assistant_callback.assert_not_called()
agent._execute_tool_calls.assert_not_called()
assert agent.client is not None
assert agent.client.chat.completions.create.call_count == 1
assert result["failed"] is True
assert result["completed"] is False
assert result["turn_exit_reason"] == "session_persistence_failed"
# ---------------------------------------------------------------------------
# Contract 2: the SEQUENTIAL path flushes each tool result immediately, BEFORE
# the next tool dispatches. Dispatch goes through run_agent.handle_function_call
@ -198,6 +300,145 @@ def test_execute_tool_calls_sequential_flushes_each_tool_result_before_next_disp
]
@pytest.mark.parametrize("executor_mode", ["sequential", "concurrent"])
def test_tool_result_is_durable_before_ui_completion_on_abnormal_exit(
tmp_path,
executor_mode,
):
"""A visible tool completion must already exist in the canonical DB."""
agent = _make_agent()
db_path = tmp_path / "state.db"
session_id = f"tool-result-abnormal-exit-{executor_mode}"
db = _attach_real_session_db(agent, db_path, session_id)
tool_call = _mock_tool_call(call_id="visible-call")
messages = [
{"role": "user", "content": "inspect the repository"},
{
"role": "assistant",
"content": "I'll inspect the repository now.",
"tool_calls": [
{
"id": "visible-call",
"type": "function",
"function": {"name": "web_search", "arguments": "{}"},
}
],
},
]
agent._flush_messages_to_session_db(messages)
roles_seen_by_ui: list[str] = []
def _ui_completion(*_args):
roles_seen_by_ui.extend(_durable_roles(db_path, session_id))
raise GeneratorExit("simulated process termination after tool completion")
agent.tool_complete_callback = _ui_completion
assistant_message = SimpleNamespace(content="", tool_calls=[tool_call])
dispatch_patch = (
patch("run_agent.handle_function_call", return_value="repository result")
if executor_mode == "sequential"
else patch.object(agent, "_invoke_tool", return_value="repository result")
)
try:
with (
dispatch_patch,
patch(
"agent.tool_executor.maybe_persist_tool_result",
side_effect=lambda **kwargs: kwargs["content"],
),
pytest.raises(GeneratorExit, match="simulated process termination"),
):
if executor_mode == "sequential":
agent._execute_tool_calls_sequential(
assistant_message,
messages,
"task-1",
)
else:
agent._execute_tool_calls_concurrent(
assistant_message,
messages,
"task-1",
)
finally:
db.close()
expected_roles = ["user", "assistant", "tool"]
assert roles_seen_by_ui == expected_roles
durable = _durable_messages(db_path, session_id)
assert [message["role"] for message in durable] == expected_roles
assert durable[2]["tool_call_id"] == "visible-call"
assert durable[2]["content"] == "repository result"
@pytest.mark.parametrize("executor_mode", ["sequential", "concurrent"])
def test_failed_tool_result_persist_blocks_completion_projection(executor_mode):
agent = _make_agent()
tool_call = _mock_tool_call(call_id="failed-persist")
assistant_message = SimpleNamespace(content="", tool_calls=[tool_call])
messages: list = []
agent._flush_messages_to_session_db = MagicMock(return_value=False)
agent.tool_complete_callback = MagicMock()
dispatch_patch = (
patch("run_agent.handle_function_call", return_value="repository result")
if executor_mode == "sequential"
else patch.object(agent, "_invoke_tool", return_value="repository result")
)
with (
dispatch_patch,
patch(
"agent.tool_executor.maybe_persist_tool_result",
side_effect=lambda **kwargs: kwargs["content"],
),
):
if executor_mode == "sequential":
agent._execute_tool_calls_sequential(
assistant_message,
messages,
"task-1",
)
else:
agent._execute_tool_calls_concurrent(
assistant_message,
messages,
"task-1",
)
agent.tool_complete_callback.assert_not_called()
assert getattr(agent, "_incremental_persistence_failed", False) is True
def test_segmented_batch_stops_before_later_segment_after_persist_failure():
agent = _make_agent()
first = _mock_tool_call(call_id="first")
second = _mock_tool_call(call_id="second")
assistant_message = SimpleNamespace(tool_calls=[first, second])
messages: list = []
agent._flush_messages_to_session_db = MagicMock(return_value=False)
with (
patch.object(agent, "_invoke_tool", return_value="first result") as invoke,
patch("run_agent.handle_function_call", return_value="second result") as dispatch,
patch(
"agent.tool_executor.maybe_persist_tool_result",
side_effect=lambda **kwargs: kwargs["content"],
),
):
execute_tool_calls_segmented(
agent,
assistant_message,
messages,
"task-1",
segments=[("parallel", [first]), ("sequential", [second])],
)
invoke.assert_called_once()
dispatch.assert_not_called()
assert getattr(agent, "_incremental_persistence_failed", False) is True
# ---------------------------------------------------------------------------
# Contract 3: the CONCURRENT path flushes each collected tool result in append
# order. Dispatch goes through agent._invoke_tool (the real concurrent