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/services/search/cache.py b/services/search/cache.py index 222682c7b..a3134abf0 100644 --- a/services/search/cache.py +++ b/services/search/cache.py @@ -40,24 +40,37 @@ def generate_cache_key(data: str) -> str: def cleanup_cache(cache_dir: Path, cache_index: Dict[str, datetime], max_age: timedelta): """Remove expired cache entries and enforce LRU policy.""" current_time = datetime.now() - files_in_dir = {f.name.split(".")[0]: f for f in cache_dir.glob("*.cache")} + files_in_dir = {f.stem: f for f in cache_dir.glob("*.cache")} - to_remove = [] - for key, timestamp in list(cache_index.items()): - if current_time - timestamp > max_age or key not in files_in_dir: - to_remove.append(key) - if key in files_in_dir: - files_in_dir[key].unlink(missing_ok=True) + live_entries = [] + for key, cache_file in list(files_in_dir.items()): + timestamp = cache_index.get(key) + if timestamp is None: + try: + timestamp = datetime.fromtimestamp(cache_file.stat().st_mtime) + except OSError: + continue + if current_time - timestamp > max_age: + try: + cache_file.unlink(missing_ok=True) + cache_metrics["evictions"] += 1 + cache_index.pop(key, None) + except OSError as e: + logger.debug("Failed to remove expired cache file %s: %s", cache_file, e) + else: + live_entries.append((key, timestamp, cache_file)) - for key in to_remove: - cache_index.pop(key, None) - cache_metrics["evictions"] += 1 - - if len(cache_index) > CACHE_MAX_ENTRIES: - sorted_items = sorted(cache_index.items(), key=lambda x: x[1]) - excess_count = len(cache_index) - CACHE_MAX_ENTRIES - for key, _ in sorted_items[:excess_count]: + for key in list(cache_index): + if key not in files_in_dir: cache_index.pop(key, None) - cache_file = cache_dir / f"{key}.cache" - cache_file.unlink(missing_ok=True) cache_metrics["evictions"] += 1 + + if len(live_entries) > CACHE_MAX_ENTRIES: + excess_count = len(live_entries) - CACHE_MAX_ENTRIES + for key, _, cache_file in sorted(live_entries, key=lambda x: x[1])[:excess_count]: + try: + cache_file.unlink(missing_ok=True) + cache_metrics["evictions"] += 1 + cache_index.pop(key, None) + except OSError as e: + logger.debug("Failed to remove excess cache file %s: %s", cache_file, e) 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 ea2d8c1bb..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); @@ -4917,7 +4920,7 @@ import { wireArrowUpRecall, getUserMessagesFromChatHistory } from './composerArr if (replaceFromHere) { // Regenerate flows intentionally trim history to this point before // resubmitting. The plain "Resend message" action must not do this. - const keepCount = msgIndex; + const keepCount = msgIndex + 1; await fetch(`${API_BASE}/api/session/${sessionId}/truncate`, { method: 'POST', headers: { 'Content-Type': 'application/json' }, @@ -4934,6 +4937,7 @@ import { wireArrowUpRecall, getUserMessagesFromChatHistory } from './composerArr sibling = next; } _hideUserBubble = true; + _pendingRegenerateSend = true; } _pendingRegenAttachments = _ids; @@ -5031,7 +5035,7 @@ import { wireArrowUpRecall, getUserMessagesFromChatHistory } from './composerArr variants.push({ raw: oldRaw, html: oldHtml, label: 'original' }); } - const keepCount = userIndex; + const keepCount = userIndex + 1; try { await fetch(`${API_BASE}/api/session/${sessionId}/truncate`, { @@ -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 c107e84fc..187411485 100644 --- a/tests/test_resend_message_nondestructive.py +++ b/tests/test_resend_message_nondestructive.py @@ -41,3 +41,18 @@ def test_only_regenerate_callers_opt_into_replace_from_here(): assert "window.chatModule.resendUserMessage(msgElement);" in renderer assert "window.chatModule.resendUserMessage(userMsgEl, { replaceFromHere: true });" in renderer + + +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 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 diff --git a/tests/test_search_cache_cleanup.py b/tests/test_search_cache_cleanup.py new file mode 100644 index 000000000..286eb8295 --- /dev/null +++ b/tests/test_search_cache_cleanup.py @@ -0,0 +1,69 @@ +import os +import time +from datetime import datetime, timedelta +from pathlib import Path + +from services.search import cache as cache_module + + +def _write_cache_file(cache_dir, key, age_seconds): + path = cache_dir / f"{key}.cache" + path.write_text("{}", encoding="utf-8") + mtime = time.time() - age_seconds + os.utime(path, (mtime, mtime)) + return path + + +def test_cleanup_cache_removes_expired_disk_files_missing_from_index(tmp_path): + expired = _write_cache_file(tmp_path, "expired", age_seconds=7200) + + cache_module.cleanup_cache(tmp_path, {}, timedelta(hours=1)) + + assert not expired.exists() + + +def test_cleanup_cache_keeps_fresh_disk_files_missing_from_index(tmp_path): + fresh = _write_cache_file(tmp_path, "fresh", age_seconds=60) + + cache_module.cleanup_cache(tmp_path, {}, timedelta(hours=1)) + + assert fresh.exists() + + +def test_cleanup_cache_enforces_max_entries_against_disk_files(tmp_path, monkeypatch): + monkeypatch.setattr(cache_module, "CACHE_MAX_ENTRIES", 2) + oldest = _write_cache_file(tmp_path, "oldest", age_seconds=30) + newer = _write_cache_file(tmp_path, "newer", age_seconds=20) + newest = _write_cache_file(tmp_path, "newest", age_seconds=10) + + cache_module.cleanup_cache(tmp_path, {}, timedelta(hours=1)) + + assert not oldest.exists() + assert newer.exists() + assert newest.exists() + + +def test_cleanup_cache_removes_index_entries_for_missing_files(tmp_path): + cache_index = {"missing": datetime.now()} + + cache_module.cleanup_cache(tmp_path, cache_index, timedelta(hours=1)) + + assert cache_index == {} + + +def test_cleanup_cache_keeps_index_when_delete_fails(tmp_path, monkeypatch): + expired = _write_cache_file(tmp_path, "expired", age_seconds=7200) + cache_index = {"expired": datetime.fromtimestamp(expired.stat().st_mtime)} + original_unlink = Path.unlink + + def fail_expired_unlink(self, missing_ok=False): + if self == expired: + raise OSError("delete failed") + return original_unlink(self, missing_ok=missing_ok) + + monkeypatch.setattr(Path, "unlink", fail_expired_unlink) + + cache_module.cleanup_cache(tmp_path, cache_index, timedelta(hours=1)) + + assert expired.exists() + assert "expired" in cache_index