This commit is contained in:
adabarbulescu 2026-08-04 02:45:47 +03:00 committed by GitHub
commit b2ee646193
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 386 additions and 21 deletions

View file

@ -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

View file

@ -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

View file

@ -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)

View file

@ -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

View file

@ -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.

View file

@ -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,
)

View file

@ -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")

View file

@ -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

View file

@ -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