From fd4c833511ac347be287a590f3e8455bd64ca670 Mon Sep 17 00:00:00 2001 From: weijia <126332148+wjl2023@users.noreply.github.com> Date: Thu, 23 Apr 2026 16:36:41 +0800 Subject: [PATCH] feat(stapp): restore chat history bubbles after /continue N (#138) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- frontends/continue_cmd.py | 51 +++++++++++++++++++++++++++++++++++++++ frontends/stapp.py | 19 +++++++++++---- 2 files changed, 65 insertions(+), 5 deletions(-) diff --git a/frontends/continue_cmd.py b/frontends/continue_cmd.py index fc8fa59..bb3c867 100644 --- a/frontends/continue_cmd.py +++ b/frontends/continue_cmd.py @@ -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() diff --git a/frontends/stapp.py b/frontends/stapp.py index b9afc4b..e5ee194 100644 --- a/frontends/stapp.py +++ b/frontends/stapp.py @@ -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')