fix(#5564): mint stable per-message ids so large-session forks align
WebUI keeps two parallel arrays per session: messages (display transcript) and context_messages (what's sent to the model). In large/compacted sessions they diverge. Forking / in-place truncation copies a prefix of both and relies on truncate_context_for_display_keep() to translate a display keep-index into the matching context index — an aligner that prefers a stable per-row `id`. But the model-context rows carried neither id nor timestamp (a 989-session live scan found 0 with any id), so alignment fell back to fragile content-signature matching that goes ambiguous on repeated tool calls / empty assistant turns. This closes the DATA gap (the alignment logic itself already prefers `id` and handles the compacted case via the matcher, shipped in #5563): mint a monotonic, session-unique integer `id` on the per-turn result rows AFTER the context restore and BEFORE both arrays are built, so the display and model-context copies of a logical row share the same id. - api/streaming.py: new _assign_stable_message_ids(); _restore_reasoning_metadata carries `id` forward across turns (as it already does timestamp); wired into all three streaming commit sites (main turn, retry, self-heal). - api/routes.py: same mint on the runs/MoA _handle_chat_sync commit path. - api/gateway_chat.py: mints ids on the two new rows of a gateway turn. - The `id` is stripped before the provider API call (not in _API_SAFE_MSG_KEYS), same treatment as timestamp — nothing new reaches the provider. session_ops.py alignment was already delivered by #5563 (id-preferring matcher + compacted-case fall-through), so no change is needed there now; tests updated to assert the post-#5563 behavior (id-bearing = exact cut; id-less = errs toward under-keeping, never the old raw-index mis-cut). Co-authored-by: b3nw <b3nw@users.noreply.github.com>
This commit is contained in:
nesquena-hermes
parent
b504fe3057
commit
ee93425b26
5 changed files with 281 additions and 2 deletions
|
|
@ -843,6 +843,19 @@ def _run_gateway_chat_streaming(
|
|||
if saved_reasoning:
|
||||
assistant_msg["reasoning"] = saved_reasoning
|
||||
previous_context = list(getattr(s, "context_messages", None) or getattr(s, "messages", None) or [])
|
||||
# Stamp stable ids on the two new rows (shared with the display merge
|
||||
# below) so display and model-context copies share an id for the
|
||||
# fork/truncate aligner (#context-message-stable-id).
|
||||
try:
|
||||
from api.streaming import _assign_stable_message_ids
|
||||
|
||||
_assign_stable_message_ids(
|
||||
[user_msg, assistant_msg],
|
||||
previous_context,
|
||||
list(getattr(s, "messages", None) or []),
|
||||
)
|
||||
except Exception:
|
||||
logger.debug("Failed to stamp stable ids on gateway turn rows", exc_info=True)
|
||||
s.context_messages = previous_context + [user_msg, assistant_msg]
|
||||
try:
|
||||
from api.streaming import _is_context_compression_marker
|
||||
|
|
|
|||
|
|
@ -20282,6 +20282,7 @@ def _handle_chat_sync(handler, body):
|
|||
)
|
||||
from api.streaming import (
|
||||
_WEBUI_PROGRESS_PROMPT,
|
||||
_assign_stable_message_ids,
|
||||
_dedupe_replayed_context_messages,
|
||||
_merge_display_messages_after_agent_result,
|
||||
_restore_display_reasoning_metadata,
|
||||
|
|
@ -20351,6 +20352,9 @@ def _handle_chat_sync(handler, body):
|
|||
_next_context_messages,
|
||||
msg,
|
||||
)
|
||||
_assign_stable_message_ids(
|
||||
_result_messages, _previous_messages, _previous_context_messages
|
||||
)
|
||||
s.context_messages = _next_context_messages
|
||||
s.messages = _merge_display_messages_after_agent_result(
|
||||
_previous_messages,
|
||||
|
|
|
|||
|
|
@ -4170,6 +4170,46 @@ def _deduplicate_context_messages(messages):
|
|||
return deduped
|
||||
|
||||
|
||||
def _assign_stable_message_ids(result_messages, *existing_arrays):
|
||||
"""Mint a stable, session-unique integer ``id`` on model-result rows lacking one.
|
||||
|
||||
Both ``messages`` (display transcript) and ``context_messages`` (model-facing
|
||||
history) are derived from the *same* per-turn ``result['messages']`` dicts.
|
||||
Stamping the id on those shared dicts here makes each logical row carry an
|
||||
identical id in both arrays, so the fork/truncate aligner
|
||||
(``session_ops.truncate_context_for_display_keep``) can match rows by its
|
||||
preferred ``id`` key instead of the fragile content-signature fallback that
|
||||
goes blind on large sessions full of structurally-identical rows.
|
||||
|
||||
Ids are monotonic within a session (max existing id + 1, seeded from the
|
||||
result rows plus any ``existing_arrays`` passed for collision-avoidance).
|
||||
Across sessions ids may repeat, which is safe: alignment only ever compares
|
||||
rows within one session, and a fork copies both arrays together so the child
|
||||
stays internally consistent. Rows whose historical id was carried forward by
|
||||
``_restore_reasoning_metadata`` keep it; only genuinely new rows are minted.
|
||||
|
||||
Returns the number of rows newly stamped. Mutates ``result_messages`` in place.
|
||||
"""
|
||||
if not result_messages:
|
||||
return 0
|
||||
seed = 0
|
||||
for arr in (result_messages, *existing_arrays):
|
||||
for m in arr or []:
|
||||
if isinstance(m, dict):
|
||||
mid = m.get('id')
|
||||
# bool is an int subclass; exclude it so a stray True/False id
|
||||
# can never seed the counter.
|
||||
if isinstance(mid, int) and not isinstance(mid, bool) and mid > seed:
|
||||
seed = mid
|
||||
stamped = 0
|
||||
for m in result_messages:
|
||||
if isinstance(m, dict) and m.get('id') is None:
|
||||
seed += 1
|
||||
m['id'] = seed
|
||||
stamped += 1
|
||||
return stamped
|
||||
|
||||
|
||||
def _prune_context_tool_results_after_compression(agent, context_messages):
|
||||
"""Run the active compressor's cheap tool-result pruning on model context.
|
||||
|
||||
|
|
@ -4235,6 +4275,13 @@ def _restore_reasoning_metadata(previous_messages, updated_messages):
|
|||
if isinstance(prev_msg, dict) and isinstance(cur_msg, dict) and _safe_projection(prev_msg) == _safe_projection(cur_msg):
|
||||
if prev_msg.get('role') == 'assistant' and prev_msg.get('reasoning') and not cur_msg.get('reasoning'):
|
||||
cur_msg['reasoning'] = prev_msg['reasoning']
|
||||
# Carry the stable per-message id (#context-message-stable-id) forward
|
||||
# the same way timestamp is carried. The agent rebuilds result rows
|
||||
# without our id every turn; without this, historical context rows
|
||||
# would be re-minted a fresh id each turn and drift out of alignment
|
||||
# with their display counterpart.
|
||||
if prev_msg.get('id') is not None and cur_msg.get('id') is None:
|
||||
cur_msg['id'] = prev_msg['id']
|
||||
if prev_msg.get('timestamp') and not cur_msg.get('timestamp'):
|
||||
cur_msg['timestamp'] = prev_msg['timestamp']
|
||||
elif prev_msg.get('_ts') and not cur_msg.get('_ts') and not cur_msg.get('timestamp'):
|
||||
|
|
@ -8186,6 +8233,13 @@ def _run_agent_streaming(
|
|||
_next_context_messages,
|
||||
msg_text,
|
||||
)
|
||||
# Stamp stable ids on the shared result rows AFTER the context
|
||||
# restore (so carried-forward ids survive) and BEFORE both
|
||||
# arrays are built, so display and model-context copies of each
|
||||
# row share an id for the fork/truncate aligner.
|
||||
_assign_stable_message_ids(
|
||||
_result_messages, _previous_messages, _previous_context_messages
|
||||
)
|
||||
s.context_messages = _deduplicate_context_messages(_next_context_messages)
|
||||
s.messages = _merge_display_messages_after_agent_result(
|
||||
_previous_messages,
|
||||
|
|
@ -8496,6 +8550,9 @@ def _run_agent_streaming(
|
|||
_next_context_messages,
|
||||
msg_text,
|
||||
)
|
||||
_assign_stable_message_ids(
|
||||
_result_messages, _previous_messages, _previous_context_messages
|
||||
)
|
||||
s.context_messages = _deduplicate_context_messages(_next_context_messages)
|
||||
s.messages = _merge_display_messages_after_agent_result(
|
||||
_previous_messages,
|
||||
|
|
@ -9552,6 +9609,9 @@ def _run_agent_streaming(
|
|||
_next_context_messages,
|
||||
msg_text,
|
||||
)
|
||||
_assign_stable_message_ids(
|
||||
_result_messages, _previous_messages, _previous_context_messages
|
||||
)
|
||||
s.context_messages = _deduplicate_context_messages(_next_context_messages)
|
||||
s.messages = _merge_display_messages_after_agent_result(
|
||||
_previous_messages,
|
||||
|
|
|
|||
183
tests/test_context_message_stable_ids.py
Normal file
183
tests/test_context_message_stable_ids.py
Normal file
|
|
@ -0,0 +1,183 @@
|
|||
"""Stable per-message ids keep the fork/truncate aligner from going blind.
|
||||
|
||||
Regression coverage for the systemic gap documented in
|
||||
local-ops/docs/bug-context-messages-lost-timestamps.md: ``context_messages``
|
||||
rows carried neither a stable ``id`` nor a ``timestamp``, so
|
||||
``truncate_context_for_display_keep`` degraded to fragile content-signature
|
||||
matching and mis-cut large sessions on fork/truncate.
|
||||
|
||||
The fix mints a stable, session-unique integer ``id`` on the per-turn result
|
||||
rows that both ``messages`` (display) and ``context_messages`` (model) derive
|
||||
from, and carries it forward across turns. These tests prove:
|
||||
|
||||
1. the minting helper is monotonic and non-destructive,
|
||||
2. ``_restore_reasoning_metadata`` carries ids forward like timestamps,
|
||||
3. the real turn transforms give a logical row the SAME id in both arrays,
|
||||
4. the aligner resolves an otherwise-ambiguous boundary via id alone.
|
||||
"""
|
||||
|
||||
from api.streaming import (
|
||||
_assign_stable_message_ids,
|
||||
_restore_reasoning_metadata,
|
||||
_deduplicate_context_messages,
|
||||
_dedupe_replayed_context_messages,
|
||||
_merge_display_messages_after_agent_result,
|
||||
)
|
||||
from api.session_ops import truncate_context_for_display_keep
|
||||
|
||||
|
||||
def test_assign_stable_ids_mints_monotonic_ints():
|
||||
rows = [
|
||||
{"role": "user", "content": "u1"},
|
||||
{"role": "assistant", "content": "a1"},
|
||||
{"role": "tool", "content": "t1", "tool_call_id": "c1"},
|
||||
]
|
||||
stamped = _assign_stable_message_ids(rows)
|
||||
assert stamped == 3
|
||||
assert [r["id"] for r in rows] == [1, 2, 3]
|
||||
|
||||
|
||||
def test_assign_stable_ids_seeds_from_existing_max_and_skips_present():
|
||||
prev = [{"role": "user", "content": "old", "id": 41}]
|
||||
rows = [
|
||||
{"role": "assistant", "content": "keep", "id": 7}, # already has id -> untouched
|
||||
{"role": "tool", "content": "new", "tool_call_id": "c"},
|
||||
]
|
||||
stamped = _assign_stable_message_ids(rows, prev)
|
||||
assert stamped == 1
|
||||
# kept its own id; new row minted above the global max (41), not above 7
|
||||
assert rows[0]["id"] == 7
|
||||
assert rows[1]["id"] == 42
|
||||
|
||||
|
||||
def test_assign_stable_ids_ignores_bool_and_empty():
|
||||
assert _assign_stable_message_ids([]) == 0
|
||||
assert _assign_stable_message_ids(None) == 0
|
||||
rows = [{"role": "user", "content": "u", "id": True}] # bool must not seed
|
||||
other = [{"role": "assistant", "content": "a"}]
|
||||
_assign_stable_message_ids(other, rows)
|
||||
assert other[0]["id"] == 1 # seed ignored the bool, started at 1
|
||||
|
||||
|
||||
def test_restore_reasoning_metadata_carries_id_forward():
|
||||
previous = [
|
||||
{"role": "user", "content": "u1", "id": 10, "timestamp": 1},
|
||||
{"role": "assistant", "content": "a1", "id": 11, "timestamp": 2},
|
||||
]
|
||||
# Agent rebuilds the same history without our id (and without timestamp).
|
||||
updated = [
|
||||
{"role": "user", "content": "u1"},
|
||||
{"role": "assistant", "content": "a1"},
|
||||
{"role": "user", "content": "u2"}, # genuinely new tail row
|
||||
]
|
||||
out = _restore_reasoning_metadata(previous, updated)
|
||||
assert out[0]["id"] == 10
|
||||
assert out[1]["id"] == 11
|
||||
assert "id" not in out[2] # new row left for the minting pass
|
||||
|
||||
|
||||
def test_both_arrays_share_id_for_same_logical_row():
|
||||
"""The real turn transforms must give a row the SAME id in both arrays."""
|
||||
prev_display = []
|
||||
prev_context = []
|
||||
# The two arrays are built from the SAME result dicts (as in streaming.py).
|
||||
result_messages = [
|
||||
{"role": "user", "content": "hello"},
|
||||
{"role": "assistant", "content": "", "tool_calls": [{"id": "c1", "name": "x"}]},
|
||||
{"role": "tool", "content": "out", "tool_call_id": "c1"},
|
||||
{"role": "assistant", "content": "done"},
|
||||
]
|
||||
|
||||
# --- context path (mirrors streaming.py commit site) ---
|
||||
next_ctx = _restore_reasoning_metadata(prev_context, result_messages)
|
||||
next_ctx = _dedupe_replayed_context_messages(prev_context, next_ctx, "hello")
|
||||
_assign_stable_message_ids(result_messages, prev_display, prev_context)
|
||||
context_out = _deduplicate_context_messages(next_ctx)
|
||||
|
||||
# --- display path (mirrors streaming.py commit site) ---
|
||||
display_out = _merge_display_messages_after_agent_result(
|
||||
prev_display,
|
||||
prev_context,
|
||||
_restore_reasoning_metadata(prev_display, result_messages),
|
||||
"hello",
|
||||
)
|
||||
|
||||
# every context row must have an id, and the matching display row (by
|
||||
# role+content) must carry the identical id.
|
||||
assert all(isinstance(m.get("id"), int) for m in context_out)
|
||||
disp_by_key = {(m["role"], str(m.get("content"))): m.get("id") for m in display_out}
|
||||
for m in context_out:
|
||||
assert disp_by_key.get((m["role"], str(m.get("content")))) == m["id"]
|
||||
|
||||
|
||||
def _dup_ambiguous_shapes(with_ids):
|
||||
"""Large-session shape: display longer than context, full of structurally
|
||||
identical assistant(tool_call)/tool(result) pairs so content signatures
|
||||
collide. Context is the compressed/shorter model view.
|
||||
"""
|
||||
dup_call = [{"id": "c", "name": "run"}]
|
||||
|
||||
def row(role, content, **kw):
|
||||
r = {"role": role, "content": content}
|
||||
r.update(kw)
|
||||
return r
|
||||
|
||||
def maybe_id(r, i):
|
||||
if with_ids and i is not None:
|
||||
r["id"] = i
|
||||
return r
|
||||
|
||||
# 8 display rows, three identical assistant/tool pairs then a trailing user.
|
||||
display = [
|
||||
maybe_id(row("user", "u1"), 1),
|
||||
maybe_id(row("assistant", "", tool_calls=dup_call), 2),
|
||||
maybe_id(row("tool", "same", tool_call_id="c"), 3),
|
||||
maybe_id(row("assistant", "", tool_calls=dup_call), 4),
|
||||
maybe_id(row("tool", "same", tool_call_id="c"), 5),
|
||||
maybe_id(row("assistant", "", tool_calls=dup_call), 6),
|
||||
maybe_id(row("tool", "same", tool_call_id="c"), 7),
|
||||
maybe_id(row("user", "u2"), 8),
|
||||
]
|
||||
# 7 context rows (shorter -> len(ctx) < len(msgs), the documented signature):
|
||||
# a leading compaction summary with no display match, then display ids 1..6.
|
||||
context = [
|
||||
row("assistant", "[compaction summary]"), # leading, unmatched, no id
|
||||
maybe_id(row("user", "u1"), 1),
|
||||
maybe_id(row("assistant", "", tool_calls=dup_call), 2),
|
||||
maybe_id(row("tool", "same", tool_call_id="c"), 3),
|
||||
maybe_id(row("assistant", "", tool_calls=dup_call), 4),
|
||||
maybe_id(row("tool", "same", tool_call_id="c"), 5),
|
||||
maybe_id(row("assistant", "", tool_calls=dup_call), 6),
|
||||
]
|
||||
return display, context
|
||||
|
||||
|
||||
def test_aligner_resolves_ambiguous_boundary_via_id():
|
||||
"""With shared ids the aligner cuts at the exact boundary even though the
|
||||
duplicate rows make content signatures ambiguous."""
|
||||
display, context = _dup_ambiguous_shapes(with_ids=True)
|
||||
# Keep the first 4 display rows (u1, a, t, a=id4). Correct cut keeps the
|
||||
# leading summary + ids 1..4 (5 rows), stopping before id5.
|
||||
out = truncate_context_for_display_keep(context, display, keep=4)
|
||||
assert len(out) == 5
|
||||
assert [m.get("id") for m in out] == [None, 1, 2, 3, 4]
|
||||
|
||||
|
||||
def test_aligner_without_ids_underkeeps_baseline():
|
||||
"""Same shape with ids/timestamps stripped (pre-fix live data): the duplicate
|
||||
rows defeat signature matching, so the aligner cannot resolve the exact
|
||||
boundary and errs toward UNDER-keeping (the #5563 matcher no longer raw-slices
|
||||
the compacted case). It keeps fewer rows than the id-aware path, which is why
|
||||
the shared stable id is the durable fix — with ids the same shape resolves to
|
||||
the exact 5-row cut (see test_aligner_resolves_ambiguous_boundary_via_id)."""
|
||||
display, context = _dup_ambiguous_shapes(with_ids=False)
|
||||
out = truncate_context_for_display_keep(context, display, keep=4)
|
||||
# Without ids the matcher can't disambiguate the identical assistant/tool
|
||||
# pairs, so it under-keeps (3 rows) rather than reaching the exact id-aware
|
||||
# boundary (5 rows). It never over-keeps or slices mid-turn at the raw index.
|
||||
assert len(out) < 5
|
||||
assert [m.get("content") for m in out] == [
|
||||
"[compaction summary]",
|
||||
"u1",
|
||||
"",
|
||||
]
|
||||
|
|
@ -793,11 +793,30 @@ def test_handle_chat_sync_writeback_dedupes_full_context_replay(tmp_path, monkey
|
|||
|
||||
assert handler.status == 200
|
||||
reloaded = Session.load(session.session_id)
|
||||
assert reloaded.context_messages == previous_context + [
|
||||
|
||||
# Stable per-message ids (#context-message-stable-id) are now stamped on
|
||||
# context rows; compare on the semantic fields so the dedup invariant is
|
||||
# still what is under test. Build ``expected`` from fresh literals: the id
|
||||
# stamping mutates the shared result dicts in place, and this test reuses the
|
||||
# same objects for its result and its previous_context.
|
||||
def _no_id(rows):
|
||||
return [{k: v for k, v in m.items() if k != "id"} for m in rows]
|
||||
|
||||
expected = [
|
||||
{"role": "assistant", "content": "cron banner"},
|
||||
{"role": "user", "content": "[Session Arc Summary (d1, node 39)]\n" + "old context\n" * 400},
|
||||
{"role": "assistant", "content": "previous answer"},
|
||||
{"role": "user", "content": "simple follow-up"},
|
||||
{"role": "assistant", "content": "short answer"},
|
||||
]
|
||||
assert reloaded.context_messages.count(previous_context[0]) == 1
|
||||
assert _no_id(reloaded.context_messages) == expected
|
||||
assert _no_id(reloaded.context_messages).count(expected[0]) == 1
|
||||
# The new turn's rows carry unique stable ids. (This session preloads
|
||||
# id-less legacy rows, which stay id-less until they age out of context;
|
||||
# brand-new sessions get full id coverage since previous_context is empty.)
|
||||
new_ids = [m.get("id") for m in reloaded.context_messages[-2:]]
|
||||
assert all(isinstance(i, int) for i in new_ids)
|
||||
assert len(set(new_ids)) == 2
|
||||
|
||||
|
||||
def test_session_context_falls_back_to_display_messages_for_legacy_sessions(tmp_path):
|
||||
|
|
|
|||
Loading…
Reference in a new issue