feat(stapp): restore chat history bubbles after /continue N (#138)

When the user runs '/continue N' in stapp, the agent's in-memory context
is restored, but the UI previously showed only a single ' restored' line
— all prior chat bubbles were missing.

This change parses the target session log and reconstructs the
user/assistant message pairs into st.session_state.messages, so reopening
a session feels like the conversation was never interrupted.

* continue_cmd.py: add extract_ui_messages(path)
  - parses model_responses log into [{role, content}, ...]
  - groups multi-turn LLM calls (prompts whose text starts with the
    '### [WORKING MEMORY]' header) into a single assistant bubble,
    inserting the existing '**LLM Running (Turn N) ...**' marker so
    fold_turns() renders them as collapsible segments.
  - two small helpers (_user_text / _assistant_text) keep parsing local.

* stapp.py: in the /continue branch, resolve the target log path BEFORE
  calling handle_frontend_command (which snapshots the current log and
  would otherwise shift list_sessions indices), then replace
  session_state.messages with the reconstructed history on success.
  Falls back to the previous behavior for bare /continue or failure.

Co-authored-by: wjl2023 <wjl2023@users.noreply.github.com>
This commit is contained in:
weijia
2026-04-23 16:36:41 +08:00
committed by GitHub
parent 1678114f1f
commit fd4c833511
2 changed files with 65 additions and 5 deletions

View File

@@ -214,6 +214,57 @@ def handle(agent, query, display_queue):
return query
def _user_text(prompt_body):
"""User-typed text from a prompt JSON; '' if this is an agent auto-continuation."""
try: msg = json.loads(prompt_body)
except Exception: return ''
if not isinstance(msg, dict): return ''
for blk in msg.get('content', []) or []:
if isinstance(blk, dict) and blk.get('type') == 'text':
t = (blk.get('text') or '').strip()
if t and not t.startswith('### [WORKING MEMORY]'): return t
return ''
def _assistant_text(response_body):
"""Joined text from a response blocks repr; '' on parse failure."""
try: blocks = ast.literal_eval(response_body)
except Exception: return ''
if not isinstance(blocks, list): return ''
return '\n'.join(b['text'] for b in blocks
if isinstance(b, dict) and b.get('type') == 'text'
and isinstance(b.get('text'), str) and b['text'].strip())
_TURN_MARK = '**LLM Running (Turn {}) ...**\n\n'
def extract_ui_messages(path):
"""Parse a model_responses log into [{role, content}, ...] for UI replay.
Auto-continuation turns are folded into one assistant bubble with Turn markers,
matching live chat rendering via fold_turns().
"""
try:
with open(path, encoding='utf-8', errors='replace') as f: content = f.read()
except Exception: return []
rounds = [] # [(user_text, [turn_text, ...]), ...]
for prompt, response in _pairs(content):
user = _user_text(prompt)
if user or not rounds: rounds.append((user, []))
rounds[-1][1].append(_assistant_text(response))
out = []
for user, turns in rounds:
if not user or not any(turns): continue
body = '\n\n'.join(t if i == 0 else _TURN_MARK.format(i + 1) + t
for i, t in enumerate(turns))
out += [{'role': 'user', 'content': user},
{'role': 'assistant', 'content': body}]
return out
def handle_frontend_command(agent, query, exclude_pid=None):
"""Frontend-friendly /continue entry that returns text directly."""
s = (query or '').strip()

View File

@@ -15,7 +15,7 @@ import streamlit as st
import time, json, re, threading, queue
from agentmain import GeneraticAgent
import chatapp_common # activate /continue command (monkey patches GeneraticAgent)
from continue_cmd import handle_frontend_command, reset_conversation
from continue_cmd import handle_frontend_command, reset_conversation, list_sessions, extract_ui_messages
st.set_page_config(page_title="Cowork", layout="wide")
@@ -190,10 +190,19 @@ if prompt := st.chat_input("any task?"):
st.session_state.messages = [{"role": "assistant", "content": reset_conversation(agent), "time": ts}]
_reset_and_rerun()
if cmd.startswith("/continue"):
st.session_state.messages = list(st.session_state.messages) + [
{"role": "user", "content": cmd, "time": ts},
{"role": "assistant", "content": handle_frontend_command(agent, cmd), "time": ts},
]
m = re.match(r'/continue\s+(\d+)\s*$', cmd.strip())
sessions = list_sessions(exclude_pid=os.getpid()) if m else []
idx = int(m.group(1)) - 1 if m else -1
# Resolve target path BEFORE handle (which snapshots current log, shifting indices).
target = sessions[idx][0] if 0 <= idx < len(sessions) else None
result = handle_frontend_command(agent, cmd)
history = extract_ui_messages(target) if target and result.startswith('') else None
tail = [{"role": "assistant", "content": result, "time": ts}]
if history:
st.session_state.messages = history + tail
else:
st.session_state.messages = list(st.session_state.messages) + \
[{"role": "user", "content": cmd, "time": ts}] + tail
_reset_and_rerun()
st.session_state.messages.append({"role": "user", "content": prompt})
if hasattr(agent, '_pet_req') and not prompt.startswith('/'): agent._pet_req('state=walk')