wuyang/zk-data-agent
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix runtime guidance transcript replay

Browse Source
This commit is contained in:
wuyang6
2026-06-12 19:41:26 +08:00
parent 9237d6782a
commit aac6383fe6
2 changed files with 111 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/session_store.py
+94 -5
View File
@@ -662,14 +662,43 @@ def _sync_agent_display_messages(
if str(row['message_key'] or '')
}
next_seq = max(existing.keys(), default=0) + 1
seq_offset = 0
keep_seqs: set[int] = set()
for seq, row in existing.items():
try:
existing_message = json.loads(row['message_json'])
except (TypeError, json.JSONDecodeError):
existing_message = {}
if isinstance(existing_message, dict) and _is_persistent_display_message(
existing_message
):
keep_seqs.add(seq)
for index, message in enumerate(display_messages, start=1):
if not isinstance(message, dict):
continue
message_json = json.dumps(message, ensure_ascii=False, sort_keys=True)
message_key = _display_message_key(message)
row = existing.get(index)
target_seq = index + seq_offset
row = existing.get(target_seq)
while row is not None:
try:
old_target_message = json.loads(row['message_json'])
except (TypeError, json.JSONDecodeError):
old_target_message = {}
if not isinstance(old_target_message, dict):
old_target_message = {}
if not _is_persistent_display_message(old_target_message):
break
# Display-only messages, such as runtime guidance, are inserted
# outside the model transcript. Later full-session syncs must
# preserve them and shift model messages after them; otherwise
# the UI sees an old pending prompt as a new user branch.
seq_offset += 1
target_seq = index + seq_offset
row = existing.get(target_seq)
keyed_row = existing_by_key.get(message_key)
if keyed_row is not None and keyed_row['message_json'] == message_json:
keep_seqs.add(int(keyed_row['seq']))
continue
if keyed_row is not None:
try:
@@ -695,6 +724,14 @@ def _sync_agent_display_messages(
int(keyed_row['seq']),
),
)
updated_row = {
'seq': int(keyed_row['seq']),
'message_key': message_key,
'message_json': message_json,
}
existing[int(keyed_row['seq'])] = updated_row # type: ignore[assignment]
existing_by_key[message_key] = updated_row # type: ignore[assignment]
keep_seqs.add(int(keyed_row['seq']))
continue
if row is None:
conn.execute(
@@ -706,15 +743,25 @@ def _sync_agent_display_messages(
""",
(
session_id,
index,
target_seq,
message_key,
str(message.get('role') or ''),
now,
message_json,
),
)
inserted_row = {
'seq': target_seq,
'message_key': message_key,
'message_json': message_json,
}
existing[target_seq] = inserted_row # type: ignore[assignment]
existing_by_key[message_key] = inserted_row # type: ignore[assignment]
next_seq = max(next_seq, target_seq + 1)
keep_seqs.add(target_seq)
continue
if row['message_json'] == message_json:
keep_seqs.add(target_seq)
continue
try:
old_message = json.loads(row['message_json'])
@@ -747,11 +794,14 @@ def _sync_agent_display_messages(
message_json,
),
)
existing_by_key[message_key] = {
appended_row = {
'seq': append_seq,
'message_key': message_key,
'message_json': message_json,
} # type: ignore[assignment]
}
existing[append_seq] = appended_row # type: ignore[assignment]
existing_by_key[message_key] = appended_row # type: ignore[assignment]
keep_seqs.add(append_seq)
continue
conn.execute(
"""
@@ -768,9 +818,32 @@ def _sync_agent_display_messages(
now,
message_json,
session_id,
index,
target_seq,
),
)
updated_row = {
'seq': target_seq,
'message_key': message_key,
'message_json': message_json,
}
existing[target_seq] = updated_row # type: ignore[assignment]
existing_by_key[message_key] = updated_row # type: ignore[assignment]
keep_seqs.add(target_seq)
stale_seqs = [
seq
for seq, row in existing.items()
if seq not in keep_seqs
and not _is_persistent_display_row(row)
]
if stale_seqs:
conn.execute(
f"""
delete from agent_display_messages
where session_id = ?
and seq in ({','.join('?' for _ in stale_seqs)})
""",
(session_id, *stale_seqs),
)
def _display_message_key(message: JSONDict) -> str:
@@ -795,6 +868,22 @@ def _is_replaceable_display_message(message: JSONDict) -> bool:
return metadata.get('kind') in {'run_status'}
def _is_persistent_display_message(message: JSONDict) -> bool:
metadata = message.get('metadata')
return (
isinstance(metadata, dict)
and metadata.get('kind') in {'runtime_guidance_display'}
)
def _is_persistent_display_row(row: sqlite3.Row | dict[str, Any]) -> bool:
try:
message = json.loads(row['message_json'])
except (TypeError, json.JSONDecodeError, KeyError):
return False
return isinstance(message, dict) and _is_persistent_display_message(message)
def _same_role_and_content(left: JSONDict, right: JSONDict) -> bool:
return (
str(left.get('role') or '') == str(right.get('role') or '')