diff --git a/routes/chat_helpers.py b/routes/chat_helpers.py index 22a334116..df7494a2a 100644 --- a/routes/chat_helpers.py +++ b/routes/chat_helpers.py @@ -9,7 +9,7 @@ import time from dataclasses import dataclass, field from typing import Any, Optional -from core.models import ChatMessage +from core.models import ChatMessage, get_session_manager_instance from core.database import SessionLocal from core.database import Session as DBSession, ModelEndpoint from src.llm_core import normalize_model_id @@ -481,6 +481,29 @@ def add_user_message(sess, chat_handler, preprocessed: PreprocessedMessage, inco chat_handler.update_session_name_if_needed(sess, preprocessed.text_for_context) +def replace_latest_user_message(sess, session_id: str, preprocessed: PreprocessedMessage) -> bool: + history = getattr(sess, "history", None) + if not history: + return False + latest = history[-1] + latest_role = latest.get("role") if isinstance(latest, dict) else getattr(latest, "role", None) + if latest_role != "user": + return False + user_meta = {"attachments": preprocessed.attachment_meta} if preprocessed.attachment_meta else None + replacement = ChatMessage("user", preprocessed.user_content, metadata=user_meta) + updated = list(history) + updated[-1] = replacement + manager = get_session_manager_instance() + if manager and getattr(manager, "replace_messages", None): + if manager.replace_messages(session_id, updated): + return True + logger.warning("Failed to persist regenerated user message for session %s", session_id) + return False + history[-1] = replacement + sess.message_count = len(history) + return True + + def fire_message_event(request, webhook_manager, session_id: str, sess, message: str, compare_mode: bool = False): """Fire webhook and event_bus events for a new user message.""" if webhook_manager and not compare_mode: @@ -687,6 +710,7 @@ async def build_chat_context( use_enhanced_message: bool = False, agent_mode: bool = False, allow_tool_preprocessing: bool = True, + regenerate: bool = False, ) -> ChatContext: """Build the full context (preface + messages) for an LLM call. @@ -707,12 +731,23 @@ async def build_chat_context( allow_tool_preprocessing=allow_tool_preprocessing, ) + raw_history = [] if incognito else (getattr(sess, "history", None) or []) + latest_raw = raw_history[-1] if raw_history else None + latest_raw_role = latest_raw.get("role") if isinstance(latest_raw, dict) else getattr(latest_raw, "role", None) + regenerate_reused_user = bool( + regenerate + and latest_raw_role == "user" + ) + # Add user message to history. Nobody/incognito uses a request-local # transcript store instead of session history so stale saved chats cannot # bleed into context and the turn is not persisted. if incognito: user_meta = {"attachments": preprocessed.attachment_meta} if preprocessed.attachment_meta else None _append_incognito_message(session_id, "user", preprocessed.user_content, user_meta) + elif regenerate_reused_user: + if not replace_latest_user_message(sess, session_id, preprocessed): + raise HTTPException(500, "Failed to update regenerated user message") else: add_user_message(sess, chat_handler, preprocessed, incognito=False) @@ -808,7 +843,13 @@ async def build_chat_context( # Build messages. In Nobody/incognito mode, never read saved session # history: the session id may be a temporary wrapper or, in buggy clients, a # stale normal session id. Only the ephemeral incognito transcript is safe. - messages = preface + (_incognito_messages(session_id) if incognito else sess.get_context_messages()) + context_history = _incognito_messages(session_id) if incognito else sess.get_context_messages() + if regenerate_reused_user: + context_history = list(context_history) + context_history[-1] = {"role": "user", "content": preprocessed.user_content} + if preprocessed.attachment_meta: + context_history[-1]["metadata"] = {"attachments": preprocessed.attachment_meta} + messages = preface + context_history # Current date/time — injected as a standalone *user*-role context message # placed immediately before the latest user turn, NOT folded into the diff --git a/routes/chat_routes.py b/routes/chat_routes.py index b081d5f1c..89180bd4f 100644 --- a/routes/chat_routes.py +++ b/routes/chat_routes.py @@ -600,6 +600,7 @@ def setup_chat_routes( use_research = chat_request.use_research time_filter = chat_request.time_filter preset_id = chat_request.preset_id + regenerate = bool(chat_request.regenerate) # Verify the caller owns this session before loading it. # Without this, any authenticated user can post into another user's chat. @@ -648,6 +649,7 @@ def setup_chat_routes( time_filter=time_filter, webhook_manager=webhook_manager, allow_tool_preprocessing=allow_tool_preprocessing, + regenerate=regenerate, ) # Research injection @@ -732,6 +734,7 @@ def setup_chat_routes( search_context = form_data.get("search_context") # pre-fetched web search results (compare mode) compare_mode = str(form_data.get("compare_mode", "")).lower() == "true" incognito = str(form_data.get("incognito", "")).lower() == "true" + regenerate = str(form_data.get("regenerate") or (body or {}).get("regenerate") or "").lower() == "true" plan_mode = str(form_data.get("plan_mode") or (body or {}).get("plan_mode") or "").lower() == "true" chat_mode = str(form_data.get("mode", "")).lower() # 'chat' or 'agent' # Workspace: confine the agent's file/shell tools to this folder. @@ -992,6 +995,7 @@ def setup_chat_routes( # index would be useless / unwanted noise. agent_mode=(chat_mode == "agent"), allow_tool_preprocessing=allow_tool_preprocessing, + regenerate=regenerate, ) _research_flags = {"do": do_research} # Mutable container for generator scope diff --git a/src/request_models.py b/src/request_models.py index f7755b1d4..ca6d48b0d 100644 --- a/src/request_models.py +++ b/src/request_models.py @@ -12,6 +12,7 @@ class ChatRequest(BaseModel): use_research: Optional[bool] = Field(default=False, description="Enable deep research") time_filter: Optional[str] = Field(default=None, description="Time filter for search") preset_id: Optional[str] = Field(default=None, description="Preset identifier") + regenerate: Optional[bool] = Field(default=False, description="Reuse the retained latest user turn") @field_validator('message') @classmethod diff --git a/static/js/chat.js b/static/js/chat.js index df4550d8d..54b19bffc 100644 --- a/static/js/chat.js +++ b/static/js/chat.js @@ -1424,6 +1424,8 @@ import { wireArrowUpRecall, getUserMessagesFromChatHistory } from './composerArr _displayOverride = null; const skipBubble = _hideUserBubble; _hideUserBubble = false; + const regenerateSend = _pendingRegenerateSend; + _pendingRegenerateSend = false; // Auto-recovery counter: carries across a turn's auto-continues, but resets // when the user genuinely sends a new message (so each task gets a fresh cap). // A real user turn (visible bubble) ALWAYS resets the budget — even if a @@ -1627,6 +1629,7 @@ import { wireArrowUpRecall, getUserMessagesFromChatHistory } from './composerArr const fd = new FormData(); fd.append('message', _finalMsgWithInject); fd.append('session', streamSessionId); + if (regenerateSend) fd.append('regenerate', 'true'); if (selectedRouteForSend.model) fd.append('selected_model', selectedRouteForSend.model); if (selectedRouteForSend.endpoint_url) fd.append('selected_endpoint_url', selectedRouteForSend.endpoint_url); if (selectedRouteForSend.endpoint_id) fd.append('selected_endpoint_id', selectedRouteForSend.endpoint_id); @@ -4934,6 +4937,7 @@ import { wireArrowUpRecall, getUserMessagesFromChatHistory } from './composerArr sibling = next; } _hideUserBubble = true; + _pendingRegenerateSend = true; } _pendingRegenAttachments = _ids; @@ -5053,6 +5057,7 @@ import { wireArrowUpRecall, getUserMessagesFromChatHistory } from './composerArr _hideUserBubble = true; const messageInput = uiModule.el('message'); messageInput.value = userText; + _pendingRegenerateSend = true; const submitBtn = document.querySelector('.send-btn'); if (submitBtn) submitBtn.click(); @@ -5068,6 +5073,7 @@ import { wireArrowUpRecall, getUserMessagesFromChatHistory } from './composerArr // File-ids carried over from the original user message during a regen, so // photos / OCR overrides survive into the new send. Consumed once. let _pendingRegenAttachments = null; + let _pendingRegenerateSend = false; /** * Called after streaming completes to attach variant navigation if this was a regen. diff --git a/tests/test_chat_helpers.py b/tests/test_chat_helpers.py index 16cb38c01..cb7145408 100644 --- a/tests/test_chat_helpers.py +++ b/tests/test_chat_helpers.py @@ -580,3 +580,196 @@ async def test_build_chat_context_keeps_cookie_user_owner_scope(monkeypatch): "preface_owner": "bob", "compact_owner": "bob", } + + +@pytest.mark.asyncio +async def test_build_chat_context_regenerate_reuses_latest_user_without_persisting_duplicate(monkeypatch): + async def fake_preprocess(chat_handler, message, att_ids, sess, **kwargs): + return PreprocessedMessage( + enhanced_message=message, + user_content=[{"type": "text", "text": message}], + text_for_context=message, + youtube_transcripts=[], + attachment_meta=[{"id": "file-1", "name": "scan.png"}], + ) + + monkeypatch.setattr(chat_helpers, "preprocess", fake_preprocess) + monkeypatch.setattr(chat_helpers, "extract_preset", lambda *_args, **_kwargs: PresetInfo(0.7, 1024, None, None)) + monkeypatch.setattr(chat_helpers, "load_prefs_for_user", lambda _user: {}) + monkeypatch.setattr(chat_helpers, "effective_user", lambda _request: "alice") + monkeypatch.setattr(chat_helpers, "_normalize_model_id_from_cache", lambda _sess: None) + monkeypatch.setattr(chat_helpers, "normalize_model_id", lambda *_args, **_kwargs: None) + async def fake_maybe_compact(sess, endpoint_url, model, messages, headers, owner=None): + return messages, 8192, False + + monkeypatch.setattr(chat_helpers, "maybe_compact", fake_maybe_compact) + monkeypatch.setattr(chat_helpers, "trim_for_context", lambda messages, _context_length: messages) + + adds = [] + monkeypatch.setattr(chat_helpers, "add_user_message", lambda *_args, **_kwargs: adds.append(1)) + replacements = [] + + class FakeManager: + def replace_messages(self, session_id, messages): + replacements.append((session_id, messages)) + sess.history = list(messages) + return True + + monkeypatch.setattr(chat_helpers, "get_session_manager_instance", lambda: FakeManager()) + + sess = SimpleNamespace( + endpoint_url="http://model.local/v1/chat/completions", + model="test-model", + headers={}, + history=[chat_helpers.ChatMessage("user", "stale OCR", metadata={"attachments": [{"id": "old-file"}]})], + ) + sess.get_context_messages = lambda: [m.to_dict() for m in sess.history] + + ctx = await build_chat_context( + sess=sess, + request=SimpleNamespace(), + chat_handler=SimpleNamespace(), + chat_processor=SimpleNamespace(build_context_preface=lambda **_kwargs: ([], [], [])), + message="original prompt", + session_id="session-1", + regenerate=True, + agent_mode=True, + ) + + assert adds == [] + assert replacements[0][0] == "session-1" + assert [(m.role, m.content, m.metadata) for m in sess.history] == [( + "user", + [{"type": "text", "text": "original prompt"}], + {"attachments": [{"id": "file-1", "name": "scan.png"}]}, + )] + user_messages = [m for m in ctx.messages if m.get("role") == "user"] + assert user_messages == [{ + "role": "user", + "content": [{"type": "text", "text": "original prompt"}], + "metadata": {"attachments": [{"id": "file-1", "name": "scan.png"}]}, + }] + + +@pytest.mark.asyncio +async def test_build_chat_context_regenerate_replaces_persisted_user_row(monkeypatch): + import core.database as cdb + import core.session_manager as sm + from tests.helpers.sqlite_db import make_temp_sqlite + + session_local, _engine, _tmpdb = make_temp_sqlite(cdb.Base.metadata) + monkeypatch.setattr(sm, "SessionLocal", session_local) + + manager = sm.SessionManager.__new__(sm.SessionManager) + manager.sessions = {} + manager.upload_handler = None + + sid = "regen-" + uuid.uuid4().hex[:8] + db = session_local() + try: + db.add(cdb.Session( + id=sid, + owner="alice", + name="chat", + model="test-model", + endpoint_url="http://model.local/v1/chat/completions", + archived=False, + message_count=1, + )) + db.add(cdb.ChatMessage( + id="stale-user", + session_id=sid, + role="user", + content="stale OCR", + )) + db.commit() + finally: + db.close() + + sess = manager.get_session(sid) + + async def fake_preprocess(chat_handler, message, att_ids, sess, **kwargs): + return PreprocessedMessage( + enhanced_message=message, + user_content=[{"type": "text", "text": message}], + text_for_context=message, + youtube_transcripts=[], + attachment_meta=[{"id": "file-1", "name": "scan.png"}], + ) + + async def fake_maybe_compact(sess, endpoint_url, model, messages, headers, owner=None): + return messages, 8192, False + + monkeypatch.setattr(chat_helpers, "preprocess", fake_preprocess) + monkeypatch.setattr(chat_helpers, "extract_preset", lambda *_args, **_kwargs: PresetInfo(0.7, 1024, None, None)) + monkeypatch.setattr(chat_helpers, "load_prefs_for_user", lambda _user: {}) + monkeypatch.setattr(chat_helpers, "effective_user", lambda _request: "alice") + monkeypatch.setattr(chat_helpers, "_normalize_model_id_from_cache", lambda _sess: None) + monkeypatch.setattr(chat_helpers, "normalize_model_id", lambda *_args, **_kwargs: None) + monkeypatch.setattr(chat_helpers, "maybe_compact", fake_maybe_compact) + monkeypatch.setattr(chat_helpers, "trim_for_context", lambda messages, _context_length: messages) + monkeypatch.setattr(chat_helpers, "get_session_manager_instance", lambda: manager) + + ctx = await build_chat_context( + sess=sess, + request=SimpleNamespace(), + chat_handler=SimpleNamespace(), + chat_processor=SimpleNamespace(build_context_preface=lambda **_kwargs: ([], [], [])), + message="current OCR", + session_id=sid, + regenerate=True, + agent_mode=True, + ) + manager.add_message(sid, chat_helpers.ChatMessage("assistant", "new answer")) + + db = session_local() + try: + rows = db.query(cdb.ChatMessage).filter(cdb.ChatMessage.session_id == sid).order_by(cdb.ChatMessage.timestamp).all() + assert [(row.role, row.content) for row in rows] == [ + ("user", "current OCR\n[Attachment: scan.png | id=file-1 | mime=application/octet-stream]"), + ("assistant", "new answer"), + ] + finally: + db.close() + + manager.sessions.clear() + reloaded = manager.get_session(sid) + assert [m["content"] for m in reloaded.get_context_messages()] == [ + "current OCR\n[Attachment: scan.png | id=file-1 | mime=application/octet-stream]", + "new answer", + ] + assert [m for m in ctx.messages if m.get("role") == "user"] == [{ + "role": "user", + "content": [{"type": "text", "text": "current OCR"}], + "metadata": {"attachments": [{"id": "file-1", "name": "scan.png"}]}, + }] + + +@pytest.mark.asyncio +async def test_build_chat_context_regenerate_fails_when_user_row_replacement_fails(monkeypatch): + async def fake_preprocess(chat_handler, message, att_ids, sess, **kwargs): + return PreprocessedMessage(message, message, message, [], []) + + class FakeManager: + def replace_messages(self, session_id, messages): + return False + + monkeypatch.setattr(chat_helpers, "preprocess", fake_preprocess) + monkeypatch.setattr(chat_helpers, "extract_preset", lambda *_args, **_kwargs: PresetInfo(0.7, 1024, None, None)) + monkeypatch.setattr(chat_helpers, "get_session_manager_instance", lambda: FakeManager()) + + sess = SimpleNamespace( + history=[chat_helpers.ChatMessage("user", "stale")], + get_context_messages=lambda: [{"role": "user", "content": "stale"}], + ) + + with pytest.raises(HTTPException): + await build_chat_context( + sess=sess, + request=SimpleNamespace(), + chat_handler=SimpleNamespace(), + chat_processor=SimpleNamespace(build_context_preface=lambda **_kwargs: ([], [], [])), + message="current", + session_id="session-1", + regenerate=True, + ) diff --git a/tests/test_chat_route_tool_policy.py b/tests/test_chat_route_tool_policy.py index babf85f6e..21ea53f52 100644 --- a/tests/test_chat_route_tool_policy.py +++ b/tests/test_chat_route_tool_policy.py @@ -79,6 +79,29 @@ def test_allow_web_search_reads_from_body_as_fallback(): ) +def test_regenerate_reads_from_body_as_fallback(): + """chat_stream must honor JSON regenerate requests, not just FormData.""" + source = _CHAT_ROUTES.read_text(encoding="utf-8") + tree = ast.parse(source) + + chat_stream_func = next( + node for node in ast.walk(tree) + if isinstance(node, ast.AsyncFunctionDef) and node.name == "chat_stream" + ) + + found_body_fallback = False + for node in ast.walk(chat_stream_func): + if isinstance(node, ast.Assign): + for target in node.targets: + if isinstance(target, ast.Name) and target.id == "regenerate": + src_segment = ast.get_source_segment(source, node) + if src_segment and "body" in src_segment: + found_body_fallback = True + assert found_body_fallback, ( + "regenerate assignment in chat_stream must fall back to JSON body" + ) + + def test_browser_form_followups_include_approval_and_send_phrases(): """Short approval replies after a form/browser turn must keep browser tools available.""" source = _CHAT_ROUTES.read_text(encoding="utf-8") diff --git a/tests/test_resend_message_nondestructive.py b/tests/test_resend_message_nondestructive.py index 426b7d4b2..7bb550953 100644 --- a/tests/test_resend_message_nondestructive.py +++ b/tests/test_resend_message_nondestructive.py @@ -43,18 +43,18 @@ def test_only_regenerate_callers_opt_into_replace_from_here(): assert "window.chatModule.resendUserMessage(userMsgEl, { replaceFromHere: true });" in renderer -def test_replace_and_regenerate_keep_the_original_user_message(): - chat = _CHAT_JS.read_text(encoding="utf-8") - resend = chat[ - chat.index("export async function resendUserMessage("): - chat.index("export async function regenerateFrom(") - ] - regenerate = chat[ - chat.index("export async function regenerateFrom("): - chat.index("// Pending variants from a regeneration", chat.index("export async function regenerateFrom(")) - ] +def test_regenerate_replace_paths_keep_user_row_and_mark_resubmit(): + body = _resend_body() + src = _CHAT_JS.read_text(encoding="utf-8") + regen = src[src.index("export async function regenerateFrom("):src.index("// Pending variants", src.index("export async function regenerateFrom("))] + replace_branch = body[body.index("if (replaceFromHere)"):body.index("_pendingRegenAttachments = _ids;")] + normal_resubmit = body[body.index("_pendingRegenAttachments = _ids;"):] - assert "const keepCount = msgIndex + 1;" in resend - assert "const keepCount = userIndex + 1;" in regenerate - assert "const keepCount = msgIndex;" not in resend - assert "const keepCount = userIndex;" not in regenerate + assert "const keepCount = msgIndex + 1;" in body + assert "const keepCount = userIndex + 1;" in regen + assert replace_branch.count("_pendingRegenerateSend = true;") == 1 + assert "_pendingRegenerateSend = true;" not in normal_resubmit + assert regen.count("_pendingRegenerateSend = true;") == 1 + assert "if (regenerateSend) fd.append('regenerate', 'true');" in src + assert "const keepCount = msgIndex;" not in body + assert "const keepCount = userIndex;" not in regen