diff --git a/app.py b/app.py index a89f80143..83d83b133 100644 --- a/app.py +++ b/app.py @@ -655,7 +655,12 @@ app.include_router(setup_emoji_routes()) # Sessions from routes.session_routes import setup_session_routes session_config = {"REQUEST_TIMEOUT": REQUEST_TIMEOUT, "OPENAI_API_KEY": OPENAI_API_KEY, "SESSIONS_FILE": SESSIONS_FILE} -app.include_router(setup_session_routes(session_manager, session_config, webhook_manager=webhook_manager)) +app.include_router(setup_session_routes( + session_manager, + session_config, + webhook_manager=webhook_manager, + upload_handler=upload_handler, +)) # Admin Danger Zone wipes (Settings → System → Danger Zone) from routes.admin_wipe_routes import setup_admin_wipe_routes @@ -684,7 +689,7 @@ app.include_router(setup_research_routes(research_handler, session_manager=sessi # History from routes.history.history_routes import setup_history_routes -app.include_router(setup_history_routes(session_manager)) +app.include_router(setup_history_routes(session_manager, upload_handler=upload_handler)) # Search from routes.search_routes import setup_search_routes @@ -763,7 +768,7 @@ app.include_router(setup_assistant_routes(task_scheduler)) # Calendar (CalDAV) from routes.calendar_routes import setup_calendar_routes -calendar_router = setup_calendar_routes() +calendar_router = setup_calendar_routes(upload_handler=upload_handler) app.include_router(calendar_router) # Shell (user-facing command execution) @@ -826,7 +831,7 @@ logger.info("Webhook & API token routes initialized") # Notes (Google Keep-style notes/todos) from routes.note_routes import setup_note_routes -app.include_router(setup_note_routes(task_scheduler)) +app.include_router(setup_note_routes(task_scheduler, upload_handler=upload_handler)) # Email from routes.email_routes import setup_email_routes diff --git a/core/database.py b/core/database.py index ade995871..d71b5c64c 100644 --- a/core/database.py +++ b/core/database.py @@ -1904,6 +1904,20 @@ def _migrate_chat_messages_fts(): conn = None try: conn = sqlite3.connect(db_path) + fts_content_expr_new = ( + "CASE WHEN instr(COALESCE(new.content, ''), ';base64,') > 0 " + "OR instr(COALESCE(new.content, ''), 'data:image/') > 0 " + "OR instr(COALESCE(new.content, ''), 'data:audio/') > 0 " + "THEN '[inline media omitted from search index]' " + "ELSE COALESCE(new.content, '') END" + ) + fts_content_expr_cm = ( + "CASE WHEN instr(COALESCE(cm.content, ''), ';base64,') > 0 " + "OR instr(COALESCE(cm.content, ''), 'data:image/') > 0 " + "OR instr(COALESCE(cm.content, ''), 'data:audio/') > 0 " + "THEN '[inline media omitted from search index]' " + "ELSE COALESCE(cm.content, '') END" + ) try: conn.execute("CREATE VIRTUAL TABLE IF NOT EXISTS temp._odysseus_fts5_probe USING fts5(content)") conn.execute("DROP TABLE IF EXISTS temp._odysseus_fts5_probe") @@ -1912,7 +1926,7 @@ def _migrate_chat_messages_fts(): return conn.executescript( - """ + f""" CREATE VIRTUAL TABLE IF NOT EXISTS chat_messages_fts USING fts5( content, message_id UNINDEXED, @@ -1920,10 +1934,14 @@ def _migrate_chat_messages_fts(): role UNINDEXED ); + DROP TRIGGER IF EXISTS chat_messages_fts_ai; + DROP TRIGGER IF EXISTS chat_messages_fts_ad; + DROP TRIGGER IF EXISTS chat_messages_fts_au; + CREATE TRIGGER IF NOT EXISTS chat_messages_fts_ai AFTER INSERT ON chat_messages BEGIN INSERT INTO chat_messages_fts(content, message_id, session_id, role) - VALUES (COALESCE(new.content, ''), new.id, new.session_id, new.role); + VALUES ({fts_content_expr_new}, new.id, new.session_id, new.role); END; CREATE TRIGGER IF NOT EXISTS chat_messages_fts_ad @@ -1935,14 +1953,14 @@ def _migrate_chat_messages_fts(): AFTER UPDATE ON chat_messages BEGIN DELETE FROM chat_messages_fts WHERE message_id = old.id; INSERT INTO chat_messages_fts(content, message_id, session_id, role) - VALUES (COALESCE(new.content, ''), new.id, new.session_id, new.role); + VALUES ({fts_content_expr_new}, new.id, new.session_id, new.role); END; """ ) conn.execute( - """ + f""" INSERT INTO chat_messages_fts(content, message_id, session_id, role) - SELECT COALESCE(cm.content, ''), cm.id, cm.session_id, cm.role + SELECT {fts_content_expr_cm}, cm.id, cm.session_id, cm.role FROM chat_messages cm WHERE NOT EXISTS ( SELECT 1 FROM chat_messages_fts fts @@ -1950,6 +1968,7 @@ def _migrate_chat_messages_fts(): ) """ ) + _scrub_legacy_chat_message_fts_media(conn) conn.commit() except Exception as e: logging.getLogger(__name__).warning(f"chat_messages FTS migration failed: {e}") @@ -1960,6 +1979,37 @@ def _migrate_chat_messages_fts(): pass +def _scrub_legacy_chat_message_fts_media(conn) -> None: + """Replace already-indexed inline media rows with searchable text only.""" + try: + from src.attachment_refs import search_index_text + except Exception as e: + logging.getLogger(__name__).warning(f"chat_messages FTS media scrub skipped: {e}") + return + + try: + rows = conn.execute( + """ + SELECT id, session_id, role, content + FROM chat_messages + WHERE instr(COALESCE(content, ''), ';base64,') > 0 + OR instr(COALESCE(content, ''), 'data:image/') > 0 + OR instr(COALESCE(content, ''), 'data:audio/') > 0 + """ + ).fetchall() + for message_id, session_id, role, content in rows: + conn.execute("DELETE FROM chat_messages_fts WHERE message_id = ?", (message_id,)) + conn.execute( + """ + INSERT INTO chat_messages_fts(content, message_id, session_id, role) + VALUES (?, ?, ?, ?) + """, + (search_index_text(content), message_id, session_id, role), + ) + except Exception as e: + logging.getLogger(__name__).warning(f"chat_messages FTS media scrub failed: {e}") + + def _migrate_add_email_smtp_security(): """Add explicit SMTP security mode for Proton Bridge/custom local SMTP.""" import sqlite3 diff --git a/core/session_manager.py b/core/session_manager.py index 491fbc078..f5024d212 100644 --- a/core/session_manager.py +++ b/core/session_manager.py @@ -16,6 +16,8 @@ from typing import Dict, Optional from .database import Session as DbSession, ChatMessage as DbChatMessage, Document as DbDocument, SessionLocal, utcnow_naive from .models import Session, ChatMessage +from src.attachment_refs import persistable_message_content +from src.upload_handler import reserve_message_upload_references # Re-export singleton accessors from models for convenience from .models import set_session_manager_instance, get_session_manager_instance @@ -72,6 +74,7 @@ class SessionManager: def __init__(self, sessions_file: str = None): # sessions_file kept for backward compat, not used self.sessions: Dict[str, Session] = {} + self.upload_handler = None self.load_sessions() # ------------------------------------------------------------------ @@ -230,17 +233,26 @@ class SessionManager: logger.warning("Dropping message for deleted session %s", session_id) return + missing_upload_id = reserve_message_upload_references( + getattr(self, "upload_handler", None), + getattr(db_session, "owner", None), + message.content, + message.metadata, + ) + if missing_upload_id: + raise ValueError( + f"Referenced upload is no longer available: {missing_upload_id}" + ) + msg_id = str(uuid.uuid4()) msg_time = datetime.utcnow() if message.metadata is None: message.metadata = {} message.metadata.setdefault('timestamp', _message_timestamp_iso(msg_time)) - # Multimodal content (image/audio attachments) is a list — serialize - # to JSON so the Text column can store it. On reload, _db_to_session - # detects the JSON-array prefix and parses it back. - _content = message.content - if isinstance(_content, list): - _content = json.dumps(_content) + # Multimodal content may contain provider data URLs for the live + # model call. Persist only readable text plus attachment references + # so chat_messages/FTS do not duplicate upload bytes. + _content = persistable_message_content(message.content, message.metadata) db_message = DbChatMessage( id=msg_id, session_id=session_id, @@ -322,6 +334,28 @@ class SessionManager: session = self.get_session(session_id) db = SessionLocal() try: + db_session = db.query(DbSession).filter(DbSession.id == session_id).first() + if db_session is None: + logger.warning("Cannot replace history for missing session %s", session_id) + return False + + # Reserve every incoming attachment before removing any durable + # message row. reserve_upload() shares the upload lifecycle lock + # with cleanup, so an upload cannot be deleted between this + # ownership check/access touch and the replacement transaction. + # A failed reservation must leave the existing transcript intact. + for message in messages: + missing_upload_id = reserve_message_upload_references( + getattr(self, "upload_handler", None), + getattr(db_session, "owner", None), + message.content, + message.metadata, + ) + if missing_upload_id: + raise ValueError( + f"Referenced upload is no longer available: {missing_upload_id}" + ) + db.query(DbChatMessage).filter(DbChatMessage.session_id == session_id).delete() now = datetime.now(timezone.utc) for i, message in enumerate(messages): @@ -330,15 +364,9 @@ class SessionManager: id=msg_id, session_id=session_id, role=message.role, - # Multimodal content (image/audio attachments) is a list; - # serialize to JSON so the Text column round-trips via - # _parse_msg_content. Storing the raw list let SQLAlchemy - # bind its single-quoted repr, which _parse_msg_content - # cannot parse (it looks for double-quoted "type"), so the - # attachment was destroyed on reload. Mirrors _persist_message. - content=(json.dumps(message.content) - if isinstance(message.content, list) - else message.content), + # Mirrors _persist_message: keep raw media bytes out of the + # persisted transcript and search index. + content=persistable_message_content(message.content, message.metadata), meta_data=json.dumps(message.metadata) if message.metadata else None, timestamp=now + timedelta(microseconds=i), ) @@ -347,12 +375,10 @@ class SessionManager: message.metadata = {} message.metadata["_db_id"] = msg_id - db_session = db.query(DbSession).filter(DbSession.id == session_id).first() - if db_session: - db_session.message_count = len(messages) - db_session.updated_at = now - db_session.last_accessed = now - db_session.last_message_at = now + db_session.message_count = len(messages) + db_session.updated_at = now + db_session.last_accessed = now + db_session.last_message_at = now db.commit() session.history = list(messages) diff --git a/docs/attachments.md b/docs/attachments.md new file mode 100644 index 000000000..93f9e0ffe --- /dev/null +++ b/docs/attachments.md @@ -0,0 +1,85 @@ +# Attachment References and Upload Storage + +Odysseus stores uploaded bytes once under the configured upload directory and +passes stable references through chat history, tools, and future artifact work. +The goal is to avoid duplicating large inline media payloads in +`chat_messages.content` or the SQLite FTS index. + +## Reference Shape + +Attachment references use this minimum shape: + +```json +{ + "type": "attachment_ref", + "attachment_id": "32hex-or-32hex.ext", + "name": "original-filename.png", + "mime": "image/png", + "size": 12345, + "checksum_sha256": "hex-digest", + "created_at": "2026-07-09T12:00:00" +} +``` + +Optional fields such as `width`, `height`, `vision`, `vision_model`, and +`gallery_id` may be present when the uploader or preprocessing path knows them. + +## Persistence + +The live model call may still receive provider-specific multimodal blocks for +the current turn. Persistence is different: + +- `chat_messages.content` stores readable text plus compact attachment reference + lines, never raw `data:*;base64,...` upload bytes. +- `chat_messages.metadata.attachments` stores structured attachment reference + metadata for UI reloads and future processing. +- The SQLite FTS migration recreates chat-message FTS triggers so new rows do + not index inline media payloads, and it scrubs legacy rows that were already + indexed with data URLs. + +## Tool Access + +Agent/tool context receives upload entries as `attachment_ref` manifests with an +`odysseus://attachment/` URI and `read_policy: "owner_checked_upload"`. + +For compatibility with existing built-in tools, a local `path` may be included +only after all of these checks pass: + +- the upload ID resolves through `UploadHandler.resolve_upload`; +- the requested owner is allowed to read the upload; +- the file remains inside the configured upload directory; +- the file path is inside the tool-readable roots. + +External MCP/custom tools should treat the URI and attachment ID as the stable +contract and request bytes through an owner-checked server path, not by assuming +host filesystem layout. + +## Retention and Deletion + +Current retention behavior is conservative: + +- uploads are indexed in `uploads.json` with owner, checksum, MIME type, size, + and creation time; +- admin cleanup first scans persisted chat metadata/content, document versions, + PDF source markers, gallery hashes, notes, and calendar records for live + references; +- cleanup fails closed if that reference scan cannot complete, and the lower-level + cleanup API removes nothing unless it receives a complete reference snapshot; +- expired, unreferenced uploads are removed during the completed scan, while + attachment-bearing writers must first take an owner-checked reservation that + serializes with deletion and refreshes the upload's access timestamp; +- deliberate removal atomically drops matching `uploads.json` rows before deleting + the bytes and restores those rows if filesystem removal fails; +- deleting a chat removes the chat rows but does not immediately delete shared + upload bytes, because the same upload may also be referenced by gallery items, + documents, duplicate-upload rows, or future artifact records. + +There is no distinct artifact table in the current schema. Artifact-like upload +references persisted in chat or document text are covered by the canonical +attachment-ID scan; any future artifact store must be added to reference discovery +before cleanup is allowed to consider its uploads unreferenced. + +Cleanup and write reservations share the upload-index lock. This closes the +scan/write/delete race in the documented single-worker deployment; a future +multi-process deployment must add an inter-process lock or move lifecycle state +into the database before enabling destructive cleanup in more than one worker. diff --git a/routes/calendar_routes.py b/routes/calendar_routes.py index 31efafcbc..6e0ee124c 100644 --- a/routes/calendar_routes.py +++ b/routes/calendar_routes.py @@ -13,8 +13,9 @@ from sqlalchemy import or_, and_ from dateutil.rrule import rrulestr from core.database import SessionLocal, CalendarCal, CalendarDeletedEvent, CalendarEvent -from src.auth_helpers import require_user +from src.auth_helpers import effective_user, require_user from src.upload_limits import read_upload_limited, ICS_MAX_BYTES +from src.upload_handler import reserve_upload_references logger = logging.getLogger(__name__) @@ -697,9 +698,18 @@ def _expand_rrule( # ── Routes ── -def setup_calendar_routes() -> APIRouter: +def setup_calendar_routes(upload_handler=None) -> APIRouter: router = APIRouter(prefix="/api/calendar", tags=["calendar"]) + def _reserve_calendar_uploads(request: Request, *values) -> None: + missing_id = reserve_upload_references( + upload_handler, + effective_user(request), + *values, + ) + if missing_id: + raise HTTPException(409, f"Referenced upload is no longer available: {missing_id}") + # ── CalDAV multi-account helpers ───────────────────────────────────────── def _get_caldav_accounts(owner: str) -> list: @@ -1087,6 +1097,7 @@ def setup_calendar_routes() -> APIRouter: @router.post("/events") async def create_event(request: Request, data: EventCreate): owner = _require_user(request) + _reserve_calendar_uploads(request, data.color, data.description, data.location) db = SessionLocal() try: cal = None @@ -1148,6 +1159,7 @@ def setup_calendar_routes() -> APIRouter: @router.put("/events/{uid}") async def update_event(request: Request, uid: str, data: EventUpdate): owner = _require_user(request) + _reserve_calendar_uploads(request, data.color, data.description, data.location) try: base_uid = _resolve_base_uid(uid) except ValueError as e: @@ -1241,6 +1253,7 @@ def setup_calendar_routes() -> APIRouter: @router.post("/calendars") async def create_calendar(request: Request, name: str = "Imported", color: str = "#5b8abf"): owner = _require_user(request) + _reserve_calendar_uploads(request, color) db = SessionLocal() try: cal = CalendarCal( @@ -1263,6 +1276,7 @@ def setup_calendar_routes() -> APIRouter: @router.put("/calendars/{cal_id}") async def update_calendar(request: Request, cal_id: str, name: str = None, color: str = None): owner = _require_user(request) + _reserve_calendar_uploads(request, color) db = SessionLocal() try: cal = _get_or_404_calendar(db, cal_id, owner) diff --git a/routes/chat_helpers.py b/routes/chat_helpers.py index cdd204ec2..63c8abc8f 100644 --- a/routes/chat_helpers.py +++ b/routes/chat_helpers.py @@ -17,6 +17,7 @@ from src.context_compactor import maybe_compact, trim_for_context from src.model_context import estimate_tokens from src.auth_helpers import effective_user from src.prompt_security import untrusted_context_message +from src.attachment_refs import attachment_ref from routes.prefs_routes import _load_for_user as load_prefs_for_user from fastapi import HTTPException @@ -418,13 +419,16 @@ def build_uploaded_file_manifest(att_ids: list, upload_handler, owner: Optional[ except Exception: path = None - manifest.append({ - "id": info.get("id") or str(att_id), - "name": info.get("name") or info.get("original_name") or str(att_id), - "mime": info.get("mime", ""), - "size": info.get("size", 0), + ref = attachment_ref({**info, "id": info.get("id") or str(att_id)}) + ref.update({ + "id": ref["attachment_id"], + "uri": f"odysseus://attachment/{ref['attachment_id']}", + "read_policy": "owner_checked_upload", + # Transitional compatibility: existing built-in tools can still use + # this path, but only after owner, upload-root, and tool-root checks. "path": path, }) + manifest.append(ref) return manifest diff --git a/routes/chat_routes.py b/routes/chat_routes.py index ca184c5a5..b8d9934b4 100644 --- a/routes/chat_routes.py +++ b/routes/chat_routes.py @@ -1450,7 +1450,9 @@ def setup_chat_routes( "tool_start", "tool_output", "agent_step", "doc_stream_open", "doc_stream_delta", "doc_update", "doc_suggestions", "ui_control", - "rounds_exhausted", + "rounds_exhausted", "budget_exceeded", + "loop_breaker_triggered", + "intent_nudge_exhausted", "ask_user", "plan_update", ): diff --git a/routes/document_routes.py b/routes/document_routes.py index e1b395e72..dae8b09fa 100644 --- a/routes/document_routes.py +++ b/routes/document_routes.py @@ -12,6 +12,7 @@ from core.database import SessionLocal, Document, DocumentVersion from core.database import Session as DbSession from src.auth_helpers import get_current_user, _auth_disabled from src.constants import MAIL_ATTACHMENTS_DIR +from src.upload_handler import reserve_upload_references logger = logging.getLogger(__name__) @@ -78,6 +79,14 @@ from routes.document_helpers import ( def setup_document_routes(session_manager, upload_handler=None) -> APIRouter: router = APIRouter(tags=["documents"]) + def _reserve_document_uploads(user: Optional[str], content: str) -> None: + missing_id = reserve_upload_references(upload_handler, user, content) + if missing_id: + raise HTTPException( + 409, + f"Referenced upload is no longer available: {missing_id}", + ) + def _locate_current_user_upload(request: Request, upload_id: str, user: Optional[str]): if upload_handler is None: return None @@ -124,6 +133,7 @@ def setup_document_routes(session_manager, upload_handler=None) -> APIRouter: if _looks_like_email_document(req.content, req.title): language = "email" + _reserve_document_uploads(user, req.content) _assert_pdf_marker_upload_owned(request, req.content, user, upload_handler) # Reply drafts are keyed to the source email. If a UI/tool path tries @@ -636,6 +646,7 @@ def setup_document_routes(session_manager, upload_handler=None) -> APIRouter: if doc.current_content == incoming_content and not req.force_version: return _doc_to_dict(doc) + _reserve_document_uploads(user, incoming_content) _assert_pdf_marker_upload_owned(request, incoming_content, user, upload_handler) # Check if we can coalesce with the latest version diff --git a/routes/history/history_routes.py b/routes/history/history_routes.py index 2324ae286..d0d45e8eb 100644 --- a/routes/history/history_routes.py +++ b/routes/history/history_routes.py @@ -10,7 +10,9 @@ from fastapi import APIRouter, Request, HTTPException from core.models import ChatMessage from core.database import SessionLocal, ChatMessage as DbChatMessage, Session as DbSession +from src.auth_helpers import effective_user from src.topic_analyzer import analyze_topics +from src.upload_handler import reserve_message_upload_references from routes.session_routes import ( _message_role, _message_text, @@ -98,9 +100,29 @@ def _merge_continue_rows_to_delete(db_messages, db1, db2): return to_delete -def setup_history_routes(session_manager) -> APIRouter: +def setup_history_routes(session_manager, upload_handler=None) -> APIRouter: router = APIRouter(tags=["history"]) + def _reserve_message_uploads( + request: Request, + content: Any, + metadata: Any = None, + ) -> None: + try: + missing_id = reserve_message_upload_references( + upload_handler, + effective_user(request), + content, + metadata, + ) + except (TypeError, ValueError) as exc: + raise HTTPException(400, "Invalid message attachment metadata") from exc + if missing_id: + raise HTTPException( + 409, + f"Referenced upload is no longer available: {missing_id}", + ) + def _db_history_entry(m: DbChatMessage) -> Dict[str, Any]: entry = {"role": m.role, "content": _history_display_content(m.content)} meta = {} @@ -251,7 +273,9 @@ def setup_history_routes(session_manager) -> APIRouter: content = body.get("content", "") if not content: raise HTTPException(400, "content is required") - msg = ChatMessage(role=role, content=content, metadata=body.get("metadata")) + metadata = body.get("metadata") + _reserve_message_uploads(request, content, metadata) + msg = ChatMessage(role=role, content=content, metadata=metadata) session_manager.add_message(session_id, msg) return {"status": "ok"} except KeyError: @@ -331,6 +355,8 @@ def setup_history_routes(session_manager) -> APIRouter: if not msg_id or content is None: raise HTTPException(400, "msg_id and content are required") + _reserve_message_uploads(request, content) + session = session_manager.get_session(session_id) db = SessionLocal() try: diff --git a/routes/note_routes.py b/routes/note_routes.py index ec8d14925..ec4507264 100644 --- a/routes/note_routes.py +++ b/routes/note_routes.py @@ -13,6 +13,7 @@ from core.database import SessionLocal, Note from core.middleware import INTERNAL_TOOL_USER from src.auth_helpers import require_user from src.constants import DATA_DIR +from src.upload_handler import reserve_upload_references from sqlalchemy.orm.attributes import flag_modified logger = logging.getLogger(__name__) @@ -574,7 +575,7 @@ async def dispatch_reminder( # Router factory # --------------------------------------------------------------------------- -def setup_note_routes(task_scheduler=None): +def setup_note_routes(task_scheduler=None, upload_handler=None): # Expose the scheduler to module-level `dispatch_reminder` so reminders # can also push to the in-app notification queue (the polling system # turns each entry into a real browser Notification + the existing @@ -596,6 +597,11 @@ def setup_note_routes(task_scheduler=None): # did not. return require_user(request) or None + def _reserve_note_uploads(owner: Optional[str], *values) -> None: + missing_id = reserve_upload_references(upload_handler, owner, *values) + if missing_id: + raise HTTPException(409, f"Referenced upload is no longer available: {missing_id}") + def _is_admin_or_single_user(request: Request, user: str | None) -> bool: if user == INTERNAL_TOOL_USER: return True @@ -645,6 +651,13 @@ def setup_note_routes(task_scheduler=None): @router.post("") def create_note(request: Request, body: NoteCreate): user = _owner(request) + _reserve_note_uploads( + user, + body.image_url, + body.color, + body.content, + json.dumps(body.items) if body.items is not None else None, + ) db = SessionLocal() try: note = Note( @@ -702,6 +715,13 @@ def setup_note_routes(task_scheduler=None): if user is not None and note.owner != user: raise HTTPException(404, "Note not found") + _reserve_note_uploads( + user, + body.image_url, + body.color, + body.content, + json.dumps(body.items) if body.items is not None else None, + ) if body.title is not None: note.title = body.title if body.content is not None: diff --git a/routes/session_routes.py b/routes/session_routes.py index 2d3543e36..2d8a6d87f 100644 --- a/routes/session_routes.py +++ b/routes/session_routes.py @@ -13,6 +13,7 @@ from src.request_models import SessionResponse from core.database import Session as DbSession, SessionLocal, Document, GalleryImage, utcnow_naive from src.auth_helpers import effective_user, _auth_disabled, owner_filter from src.session_actions import is_session_recently_active +from src.upload_handler import reserve_message_upload_references def _sanitize_export_filename(name: str) -> str: @@ -203,7 +204,12 @@ def _pick_endpoint_for_sort(owner=None): return url, model, headers return None, None, None -def setup_session_routes(session_manager: SessionManager, config: dict, webhook_manager=None): +def setup_session_routes( + session_manager: SessionManager, + config: dict, + webhook_manager=None, + upload_handler=None, +): """Setup session routes with the provided manager and config""" REQUEST_TIMEOUT = config.get("REQUEST_TIMEOUT", 20) @@ -537,6 +543,22 @@ def setup_session_routes(session_manager: SessionManager, config: dict, webhook_ body = await request.json() messages = body.get("messages", []) from core.models import ChatMessage + owner = effective_user(request) + try: + for message in messages: + missing_id = reserve_message_upload_references( + upload_handler, + owner, + message.get("content"), + message.get("metadata"), + ) + if missing_id: + raise HTTPException( + 409, + f"Referenced upload is no longer available: {missing_id}", + ) + except (AttributeError, TypeError, ValueError) as exc: + raise HTTPException(400, "Invalid message attachment metadata") from exc for m in messages: sess.add_message(ChatMessage(m["role"], m["content"], metadata=m.get("metadata"))) session_manager.save_sessions() diff --git a/routes/upload_routes.py b/routes/upload_routes.py index 1bd402ac8..fb702e45a 100644 --- a/routes/upload_routes.py +++ b/routes/upload_routes.py @@ -10,16 +10,140 @@ from fastapi import APIRouter, Request, File, UploadFile, HTTPException, Form from typing import List, Optional import logging from core.middleware import require_admin -from core.database import SessionLocal, GalleryImage, Session as DbSession +from core.database import ( + SessionLocal, + ChatMessage as DbChatMessage, + CalendarCal, + CalendarEvent, + Document, + DocumentVersion, + GalleryImage, + Note, + Session as DbSession, +) from src.auth_helpers import effective_user +from src.attachment_refs import attachment_refs_from_metadata from src.constants import GENERATED_IMAGES_DIR -from src.upload_handler import count_recent_uploads +from src.upload_handler import ( + UploadCleanupSafetyError, + count_recent_uploads, + extract_upload_ids, +) logger = logging.getLogger(__name__) router = APIRouter(prefix="/api/upload", tags=["upload"]) UPLOAD_RESPONSE_HEADERS = {"X-Content-Type-Options": "nosniff"} +def _upload_ids_from_persisted_text(value: object) -> set[str]: + """Return canonical upload IDs embedded in persisted text. + + This covers attachment reference lines/URIs and the PDF source markers + stored by the document editor. False positives are intentionally + conservative: retaining an extra upload is safer than deleting referenced + bytes. + """ + return extract_upload_ids(value) + + +def _upload_ids_from_message_metadata(raw_metadata: object) -> set[str]: + """Extract attachment IDs from a persisted chat metadata JSON value. + + Malformed metadata raises instead of being treated as an empty reference + set. The admin cleanup route catches that failure and aborts cleanup. + """ + if raw_metadata in (None, ""): + return set() + if isinstance(raw_metadata, str): + metadata = json.loads(raw_metadata) + else: + metadata = raw_metadata + if not isinstance(metadata, dict): + raise ValueError("chat message metadata must be a JSON object") + + attachments = metadata.get("attachments") + if attachments is not None: + if not isinstance(attachments, list) or any( + not isinstance(item, dict) for item in attachments + ): + raise ValueError("chat message attachments metadata is malformed") + + ids = { + str(ref["attachment_id"]) + for ref in attachment_refs_from_metadata(metadata) + if ref.get("attachment_id") + } + # Preserve canonical IDs even in older metadata shapes not normalized by + # attachment_refs_from_metadata(). + ids.update(_upload_ids_from_persisted_text(json.dumps(metadata))) + return ids + + +def _collect_persisted_upload_references() -> tuple[set[str], set[str]]: + """Collect upload IDs/hashes still referenced by durable application data. + + The caller must treat any exception as an incomplete scan and fail closed. + There is no distinct artifact table in the current schema; artifact-like + attachment references persisted in chat/document text are covered by the + canonical-ID scan. + """ + referenced_ids: set[str] = set() + referenced_hashes: set[str] = set() + db = SessionLocal() + try: + for content, raw_metadata in db.query( + DbChatMessage.content, + DbChatMessage.meta_data, + ).yield_per(500): + referenced_ids.update(_upload_ids_from_persisted_text(content)) + referenced_ids.update(_upload_ids_from_message_metadata(raw_metadata)) + + for (content,) in db.query(Document.current_content).yield_per(500): + referenced_ids.update(_upload_ids_from_persisted_text(content)) + + for (content,) in db.query(DocumentVersion.content).yield_per(500): + referenced_ids.update(_upload_ids_from_persisted_text(content)) + + for filename, file_hash in db.query( + GalleryImage.filename, + GalleryImage.file_hash, + ).yield_per(500): + referenced_ids.update(_upload_ids_from_persisted_text(filename)) + if file_hash: + referenced_hashes.add(str(file_hash)) + + for image_url, color, content, items in db.query( + Note.image_url, + Note.color, + Note.content, + Note.items, + ).yield_per(500): + for value in (image_url, color, content, items): + referenced_ids.update(_upload_ids_from_persisted_text(value)) + + for (color,) in db.query(CalendarCal.color).yield_per(500): + referenced_ids.update(_upload_ids_from_persisted_text(color)) + + for color, description, location in db.query( + CalendarEvent.color, + CalendarEvent.description, + CalendarEvent.location, + ).yield_per(500): + for value in (color, description, location): + referenced_ids.update(_upload_ids_from_persisted_text(value)) + + return referenced_ids, referenced_hashes + finally: + db.close() + + +def _run_reference_safe_cleanup(upload_handler) -> int: + referenced_ids, referenced_hashes = _collect_persisted_upload_references() + return upload_handler.cleanup_old_uploads( + referenced_upload_ids=referenced_ids, + referenced_upload_hashes=referenced_hashes, + ) + def setup_upload_routes(upload_handler): """Setup upload routes with the provided handler""" @@ -172,7 +296,9 @@ def setup_upload_routes(upload_handler): "mime": meta["mime"], "size": meta["size"], "hash": meta["hash"], + "checksum_sha256": meta.get("checksum_sha256") or meta["hash"], "uploaded_at": meta["uploaded_at"], + "created_at": meta.get("created_at") or meta["uploaded_at"], "width": meta.get("width"), "height": meta.get("height"), "is_duplicate": meta.get("is_duplicate", False) @@ -195,7 +321,23 @@ def setup_upload_routes(upload_handler): async def manual_cleanup(request: Request): """Manually trigger cleanup of old uploads.""" require_admin(request) - cleaned_count = upload_handler.cleanup_old_uploads() + try: + cleaned_count = await asyncio.to_thread( + _run_reference_safe_cleanup, + upload_handler, + ) + except UploadCleanupSafetyError: + logger.exception("Upload cleanup aborted because index safety checks failed") + raise HTTPException( + 503, + "Upload cleanup aborted because upload index integrity could not be verified", + ) + except Exception: + logger.exception("Upload cleanup skipped because reference discovery failed") + raise HTTPException( + 503, + "Upload cleanup skipped because persisted references could not be verified", + ) return {"status": "success", "files_cleaned": cleaned_count} @router.get("/stats") diff --git a/src/agent_loop.py b/src/agent_loop.py index 46a669c9d..6f7dde605 100644 --- a/src/agent_loop.py +++ b/src/agent_loop.py @@ -3828,9 +3828,8 @@ async def stream_agent_loop( and _intent_match is not None and len(_intent_text) < 400 and "```" not in _intent_text - and _intent_nudge_count < _MAX_INTENT_NUDGES ) - if _looks_like_promise: + if _looks_like_promise and _intent_nudge_count < _MAX_INTENT_NUDGES: _intent_nudge_count += 1 _matched_phrase = _intent_match.group(0).strip() logger.info(f"[agent] intent-without-action nudge #{_intent_nudge_count} on round {round_num}: {_matched_phrase!r}") @@ -3859,6 +3858,31 @@ async def stream_agent_loop( # Visible signal in the stream so the user knows we caught it. yield f'data: {json.dumps({"type": "agent_step", "round": round_num + 1})}\n\n' continue + if _looks_like_promise: + _matched_phrase = _intent_match.group(0).strip() + _guard_message = ( + "The agent stopped because it repeatedly announced a tool " + "action without making the tool call." + ) + logger.warning( + "[agent] intent-without-action guard exhausted on round %d after %d nudges: %r", + round_num, + _intent_nudge_count, + _matched_phrase, + ) + yield ( + "data: " + + json.dumps({ + "type": "intent_nudge_exhausted", + "reason": "intent_without_action_nudge_cap", + "message": _guard_message, + "round": round_num, + "nudges": _intent_nudge_count, + "matched": _matched_phrase, + }) + + "\n\n" + ) + break break # no tools — done # ── Loop-breaker (Terminus-style stall detector) ────────────── @@ -3895,6 +3919,21 @@ async def stream_agent_loop( reason = (f"calling {_runaway} with identical arguments over and over" if _runaway else "repeating the same tool calls without new progress") logger.warning(f"[agent] loop-breaker tripped on round {round_num} ({reason}); sig={_sig[:80]!r}") + yield ( + "data: " + + json.dumps({ + "type": "loop_breaker_triggered", + "reason": "loop_breaker_stall", + "message": ( + "The loop-breaker detected repeated tool calls without " + "new progress, so the agent is being forced to stop " + "using tools and give its best final answer." + ), + "round": round_num, + "detail": reason, + }) + + "\n\n" + ) # The model has been executing tools, so its results are already # in context. Force ONE tool-free round to converge: write the # answer from what it has, or state plainly what's blocking it. diff --git a/src/agent_tools/document_tools.py b/src/agent_tools/document_tools.py index 2ab8659d3..65ee0461e 100644 --- a/src/agent_tools/document_tools.py +++ b/src/agent_tools/document_tools.py @@ -2,10 +2,16 @@ from typing import Any, Dict, List, Optional import logging import re from src.constants import MAX_READ_CHARS -from src.tool_utils import _parse_tool_args +from src.tool_utils import _parse_tool_args, get_upload_handler +from src.upload_handler import reserve_upload_references logger = logging.getLogger(__name__) + +def _missing_document_upload(owner: Optional[str], content: Any) -> Optional[str]: + """Reserve explicit upload URLs before an agent persists document text.""" + return reserve_upload_references(get_upload_handler(), owner, content) + # --------------------------------------------------------------------------- # Active document state # --------------------------------------------------------------------------- @@ -384,6 +390,13 @@ class CreateDocumentTool: return {"error": "Cannot create document in another user's session"} _owner = _sess.owner if _sess else None + missing_id = _missing_document_upload(_owner, content) + if missing_id: + return { + "error": f"Referenced upload is no longer available: {missing_id}", + "exit_code": 1, + } + doc = Document( id=doc_id, session_id=session_id, @@ -455,6 +468,13 @@ class UpdateDocumentTool: if is_email_doc: doc.language = "email" + missing_id = _missing_document_upload(owner, new_content) + if missing_id: + return { + "error": f"Referenced upload is no longer available: {missing_id}", + "exit_code": 1, + } + if not is_email_doc and _pdf_source_upload_id(doc.current_content or ""): return _create_pdf_text_derivative( db, @@ -532,6 +552,12 @@ class EditDocumentTool: applied = 1 skipped = max(0, len(edits) - 1) doc.language = "email" + missing_id = _missing_document_upload(owner, updated_content) + if missing_id: + return { + "error": f"Referenced upload is no longer available: {missing_id}", + "exit_code": 1, + } new_ver = doc.version_count + 1 ver = DocumentVersion( id=str(uuid.uuid4()), @@ -584,6 +610,13 @@ class EditDocumentTool: if applied == 0: return {"error": f"No edits applied — none of the FIND blocks matched the document content (skipped {skipped})"} + missing_id = _missing_document_upload(owner, updated_content) + if missing_id: + return { + "error": f"Referenced upload is no longer available: {missing_id}", + "exit_code": 1, + } + if _pdf_source_upload_id(doc.current_content or ""): return _create_pdf_text_derivative( db, diff --git a/src/app_initializer.py b/src/app_initializer.py index b438b2b17..1b29f06d2 100644 --- a/src/app_initializer.py +++ b/src/app_initializer.py @@ -21,6 +21,7 @@ from src.model_discovery import ModelDiscovery from src.chat_handler import ChatHandler from src.research_handler import ResearchHandler from src.upload_handler import UploadHandler +from src.tool_utils import set_upload_handler from src.search import update_search_config logger = logging.getLogger(__name__) @@ -49,6 +50,8 @@ def initialize_managers(base_dir: str, rag_manager=None) -> Dict[str, Any]: session_manager = SessionManager(SESSIONS_FILE) set_session_manager(session_manager) # Enable Session.add_message() persistence upload_handler = UploadHandler(base_dir, UPLOAD_DIR) + session_manager.upload_handler = upload_handler + set_upload_handler(upload_handler) personal_docs_manager = PersonalDocsManager(PERSONAL_DIR, rag_manager) api_key_manager = APIKeyManager(DATA_DIR) preset_manager = PresetManager(DATA_DIR) diff --git a/src/attachment_refs.py b/src/attachment_refs.py new file mode 100644 index 000000000..054bea54b --- /dev/null +++ b/src/attachment_refs.py @@ -0,0 +1,164 @@ +"""Attachment reference helpers for chat storage and tool manifests. + +Live model calls may need provider-specific data URLs for the current turn. +Persisted history and search indexes should keep stable upload references and +human-readable text instead of duplicating raw media bytes. +""" + +from __future__ import annotations + +import json +import re +from typing import Any, Iterable + + +DATA_URL_RE = re.compile( + r"data:[^;,\s\"']+;base64,[A-Za-z0-9+/=]+", + re.IGNORECASE, +) + +MEDIA_BLOCK_TYPES = { + "image", + "image_url", + "input_image", + "audio", + "input_audio", + "file", +} + + +def strip_inline_data_urls(text: str) -> str: + """Replace inline data URLs with a compact marker.""" + if not isinstance(text, str) or ";base64," not in text: + return text + return DATA_URL_RE.sub("[inline media omitted from persisted history]", text) + + +def attachment_ref(info: dict[str, Any]) -> dict[str, Any]: + """Return the stable attachment reference shape used outside raw uploads.""" + upload_id = str(info.get("id") or info.get("attachment_id") or "").strip() + try: + size = int(info.get("size") or 0) + except (TypeError, ValueError): + size = 0 + ref = { + "type": "attachment_ref", + "attachment_id": upload_id, + "name": info.get("name") or info.get("original_name") or upload_id, + "mime": info.get("mime") or "application/octet-stream", + "size": size, + } + checksum = info.get("checksum_sha256") or info.get("sha256") or info.get("hash") + if checksum: + ref["checksum_sha256"] = checksum + created_at = info.get("created_at") or info.get("uploaded_at") + if created_at: + ref["created_at"] = created_at + for key in ("width", "height", "vision", "vision_model", "gallery_id"): + value = info.get(key) + if value is not None: + ref[key] = value + return ref + + +def attachment_refs_from_metadata(metadata: dict[str, Any] | None) -> list[dict[str, Any]]: + """Extract attachment refs from message metadata.""" + attachments = (metadata or {}).get("attachments") or [] + if not isinstance(attachments, list): + return [] + refs: list[dict[str, Any]] = [] + for item in attachments: + if isinstance(item, dict): + ref = attachment_ref(item) + if ref.get("attachment_id"): + refs.append(ref) + return refs + + +def _ref_line(ref: dict[str, Any]) -> str: + parts = [f"Attachment: {ref.get('name') or ref.get('attachment_id') or 'upload'}"] + if ref.get("attachment_id"): + parts.append(f"id={ref['attachment_id']}") + if ref.get("mime"): + parts.append(f"mime={ref['mime']}") + if ref.get("size"): + parts.append(f"size={ref['size']} bytes") + if ref.get("checksum_sha256"): + parts.append(f"sha256={ref['checksum_sha256']}") + line = "[" + " | ".join(parts) + "]" + if ref.get("vision"): + line += f"\n[Attachment description: {str(ref['vision']).strip()}]" + return line + + +def _text_from_blocks(blocks: Iterable[Any]) -> str: + lines: list[str] = [] + omitted_media = 0 + for block in blocks: + if isinstance(block, str): + lines.append(strip_inline_data_urls(block)) + continue + if not isinstance(block, dict): + continue + block_type = block.get("type") + if block_type == "text": + text = block.get("text") + if isinstance(text, str) and text: + lines.append(strip_inline_data_urls(text)) + elif block_type == "attachment_ref": + lines.append(_ref_line(block)) + elif block_type in MEDIA_BLOCK_TYPES: + omitted_media += 1 + else: + try: + encoded = json.dumps(block, ensure_ascii=True, sort_keys=True) + except TypeError: + encoded = str(block) + lines.append(strip_inline_data_urls(encoded)) + if omitted_media: + plural = "s" if omitted_media != 1 else "" + lines.append(f"[{omitted_media} inline media payload{plural} omitted]") + return "\n".join(line for line in lines if line).strip() + + +def persistable_message_content( + content: Any, + metadata: dict[str, Any] | None = None, +) -> str: + """Return content safe for DB persistence and FTS indexing. + + Multimodal provider blocks are collapsed to readable text plus stable + attachment reference lines from metadata. This avoids storing base64 media + in ``chat_messages.content`` while preserving enough context for reloads, + search, and later turns. + """ + if isinstance(content, list): + text = _text_from_blocks(content) + refs = attachment_refs_from_metadata(metadata) + ref_lines = [_ref_line(ref) for ref in refs] + if ref_lines: + text = "\n".join([part for part in (text, "\n".join(ref_lines)) if part]).strip() + return text + if isinstance(content, str): + return strip_inline_data_urls(content) + try: + return strip_inline_data_urls(json.dumps(content, ensure_ascii=True, sort_keys=True)) + except TypeError: + return strip_inline_data_urls(str(content)) + + +def search_index_text(content: Any) -> str: + """Best-effort searchable text for legacy stored content.""" + if isinstance(content, str): + raw = content.strip() + if raw.startswith("[") and '"type"' in raw: + try: + parsed = json.loads(content) + except (TypeError, ValueError): + parsed = None + if isinstance(parsed, list): + return _text_from_blocks(parsed) + return strip_inline_data_urls(content) + if isinstance(content, list): + return _text_from_blocks(content) + return persistable_message_content(content) diff --git a/src/chat_handler.py b/src/chat_handler.py index 9df6f7ce6..f9b2a2193 100644 --- a/src/chat_handler.py +++ b/src/chat_handler.py @@ -191,6 +191,8 @@ class ChatHandler: "name": fi.get("name") or fi.get("original_name") or fi["id"], "mime": fi.get("mime", ""), "size": fi.get("size", 0), + "checksum_sha256": fi.get("checksum_sha256") or fi.get("hash"), + "created_at": fi.get("created_at") or fi.get("uploaded_at"), "width": fi.get("width"), "height": fi.get("height"), }) diff --git a/src/tool_utils.py b/src/tool_utils.py index 8255bc0a9..83636cf23 100644 --- a/src/tool_utils.py +++ b/src/tool_utils.py @@ -9,6 +9,7 @@ import json from src.constants import MAX_OUTPUT_CHARS _mcp_manager = None +_upload_handler = None # --------------------------------------------------------------------------- # MCP Manager singleton @@ -23,6 +24,21 @@ def get_mcp_manager(): """Get the global MCP manager instance.""" return _mcp_manager + +# --------------------------------------------------------------------------- +# Shared upload lifecycle handler +# --------------------------------------------------------------------------- + +def set_upload_handler(handler): + """Register the process's UploadHandler without importing app modules.""" + global _upload_handler + _upload_handler = handler + + +def get_upload_handler(): + """Return the shared UploadHandler used by route and agent writers.""" + return _upload_handler + # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- diff --git a/src/tools/calendar.py b/src/tools/calendar.py index d442f2a42..e6572ba40 100644 --- a/src/tools/calendar.py +++ b/src/tools/calendar.py @@ -10,6 +10,8 @@ import re from typing import Dict, Optional from src.tools._common import _parse_tool_args +from src.tool_utils import get_upload_handler +from src.upload_handler import reserve_upload_references logger = logging.getLogger(__name__) @@ -408,11 +410,25 @@ async def do_manage_calendar(content: str, owner: Optional[str] = None) -> Dict: importance = args.get("importance") or "normal" minutes_before = _reminder_minutes(args) + event_description = _event_description(args, minutes_before) + event_location = args.get("location", "") or "" + missing_id = reserve_upload_references( + get_upload_handler(), + owner, + event_description, + event_location, + ) + if missing_id: + return { + "error": f"Referenced upload is no longer available: {missing_id}", + "exit_code": 1, + } + uid = str(_uuid.uuid4()) ev = CalendarEvent( uid=uid, calendar_id=cal.id, summary=summary, - description=_event_description(args, minutes_before), - location=args.get("location", "") or "", + description=event_description, + location=event_location, dtstart=dtstart, dtend=dtend, all_day=all_day, is_utc=dtstart_is_utc and not all_day, rrule=args.get("rrule", "") or "", @@ -465,6 +481,17 @@ async def do_manage_calendar(content: str, owner: Optional[str] = None) -> Dict: ev = _event_query().filter(CalendarEvent.uid == base_uid).first() if not ev: return {"error": f"Event {uid} not found", "exit_code": 1} + missing_id = reserve_upload_references( + get_upload_handler(), + owner, + args.get("description"), + args.get("location"), + ) + if missing_id: + return { + "error": f"Referenced upload is no longer available: {missing_id}", + "exit_code": 1, + } if args.get("summary") is not None: ev.summary = args["summary"] if args.get("description") is not None: diff --git a/src/tools/notes.py b/src/tools/notes.py index d24be55d1..ba980c158 100644 --- a/src/tools/notes.py +++ b/src/tools/notes.py @@ -10,6 +10,8 @@ import re from typing import Dict, Optional from src.tools._common import _parse_tool_args +from src.tool_utils import get_upload_handler +from src.upload_handler import reserve_upload_references logger = logging.getLogger(__name__) @@ -198,6 +200,18 @@ async def do_manage_notes(content: str, owner: Optional[str] = None) -> Dict: "duplicate": True, "exit_code": 0, } + missing_id = reserve_upload_references( + get_upload_handler(), + owner, + content_raw, + args.get("color"), + items_json, + ) + if missing_id: + return { + "error": f"Referenced upload is no longer available: {missing_id}", + "exit_code": 1, + } note = Note( id=str(_uuid.uuid4()), owner=owner, @@ -235,6 +249,19 @@ async def do_manage_notes(content: str, owner: Optional[str] = None) -> Dict: return {"error": f"Note '{note_id}' not found", "exit_code": 1} if not _note_visible_to_owner(note, owner): return {"error": "Note not found", "exit_code": 1} + missing_id = reserve_upload_references( + get_upload_handler(), + owner, + args.get("content"), + args.get("color"), + args.get("checklist_items"), + args.get("items"), + ) + if missing_id: + return { + "error": f"Referenced upload is no longer available: {missing_id}", + "exit_code": 1, + } for field in ("title", "content", "note_type", "color", "label"): if field in args and args[field] is not None: setattr(note, field, args[field]) diff --git a/src/upload_handler.py b/src/upload_handler.py index 1f24c6263..ce0b4b129 100644 --- a/src/upload_handler.py +++ b/src/upload_handler.py @@ -35,11 +35,34 @@ import logging logger = logging.getLogger(__name__) + +class UploadCleanupSafetyError(RuntimeError): + """Raised when cleanup cannot prove that destructive work is safe.""" + # The extension is optional: save_upload builds the id as `{uuid.hex}{ext}`, # and a file with no extension (Dockerfile, README, ...) yields a bare 32-hex # id. Requiring `.ext` made those ids fail validation, so the stored file # could never be resolved or downloaded again. UPLOAD_ID_RE = re.compile(r"^[0-9a-fA-F]{32}(?:\.[A-Za-z0-9]+)?$") +UPLOAD_ID_TOKEN_RE = re.compile( + r"(?\[\](){},;!?:&#]|\.(?![A-Za-z0-9]))" +) +PDF_SOURCE_UPLOAD_RE = re.compile( + r"", + re.IGNORECASE, +) +ATTACHMENT_REFERENCE_LINE_RE = re.compile( + r"\[Attachment:[^\]\r\n]*\|\s*id=" + r"([0-9a-fA-F]{32}(?:\.[A-Za-z0-9]+)?)" + r"(?:\s*\||\s*\])", + re.IGNORECASE, +) def is_valid_upload_id(upload_id: str) -> bool: @@ -47,6 +70,110 @@ def is_valid_upload_id(upload_id: str) -> bool: return UPLOAD_ID_RE.fullmatch(upload_id or "") is not None +def extract_upload_ids(value: Any) -> set[str]: + """Return canonical upload IDs embedded in a persisted URL/text value.""" + if not isinstance(value, str) or not value: + return set() + return set(UPLOAD_ID_TOKEN_RE.findall(value)) + + +def extract_internal_upload_ids(value: Any) -> set[str]: + """Return IDs from explicit internal upload references only. + + Cleanup intentionally uses :func:`extract_upload_ids` conservatively, but + write-time reservation must not treat an arbitrary 32-hex checksum in note + or calendar text as an upload reference. Nested JSON-like values are + supported because note checklist items are persisted as structured data. + """ + if isinstance(value, dict): + found: set[str] = set() + for nested in value.values(): + found.update(extract_internal_upload_ids(nested)) + return found + if isinstance(value, (list, tuple, set)): + found: set[str] = set() + for nested in value: + found.update(extract_internal_upload_ids(nested)) + return found + if not isinstance(value, str) or not value: + return set() + return ( + set(INTERNAL_UPLOAD_URL_RE.findall(value)) + | set(PDF_SOURCE_UPLOAD_RE.findall(value)) + | set(ATTACHMENT_REFERENCE_LINE_RE.findall(value)) + ) + + +def reserve_upload_references( + upload_handler: Any, + owner: Optional[str], + *values: Any, +) -> Optional[str]: + """Reserve upload IDs in values before a caller persists references. + + Returns the first ID that cannot be owner-checked/reserved, otherwise + ``None``. A missing handler is treated as no-op for backward-compatible + route factories; production wires the shared UploadHandler instance. + """ + if upload_handler is None: + return None + upload_ids: set[str] = set() + for value in values: + upload_ids.update(extract_internal_upload_ids(value)) + return reserve_upload_ids(upload_handler, owner, upload_ids) + + +def reserve_upload_ids( + upload_handler: Any, + owner: Optional[str], + upload_ids: Any, +) -> Optional[str]: + """Owner-reserve canonical IDs from a trusted structured reference field.""" + if upload_handler is None: + return None + canonical_ids = { + str(upload_id).strip() + for upload_id in (upload_ids or []) + if is_valid_upload_id(str(upload_id).strip()) + } + for upload_id in sorted(canonical_ids): + try: + resolved = upload_handler.reserve_upload( + upload_id, + owner=owner, + allow_admin=False, + ) + except Exception: + resolved = None + if not resolved: + return upload_id + return None + + +def reserve_message_upload_references( + upload_handler: Any, + owner: Optional[str], + content: Any, + metadata: Any = None, +) -> Optional[str]: + """Reserve explicit chat references, including structured attachment IDs.""" + upload_ids = extract_internal_upload_ids(content) + if metadata not in (None, ""): + if isinstance(metadata, str): + metadata = json.loads(metadata) + if not isinstance(metadata, dict): + raise ValueError("message metadata must be a JSON object") + upload_ids.update(extract_internal_upload_ids(metadata)) + from src.attachment_refs import attachment_refs_from_metadata + + upload_ids.update( + str(ref.get("attachment_id") or "").strip() + for ref in attachment_refs_from_metadata(metadata) + if ref.get("attachment_id") + ) + return reserve_upload_ids(upload_handler, owner, upload_ids) + + def _build_upload_id(safe_filename: str) -> str: """Build a unique upload id whose extension matches UPLOAD_ID_RE. @@ -249,43 +376,295 @@ class UploadHandler: return True - def cleanup_old_uploads(self): - """Remove uploaded files older than CLEANUP_DAYS days.""" + @staticmethod + def _parse_upload_timestamp(value: Any) -> Optional[datetime]: + if not isinstance(value, str) or not value.strip(): + return None try: - cutoff_date = datetime.now() - timedelta(days=self.cleanup_days) + parsed = datetime.fromisoformat(value.strip().replace("Z", "+00:00")) + if parsed.tzinfo is not None: + parsed = parsed.astimezone().replace(tzinfo=None) + return parsed + except (TypeError, ValueError): + return None + + @classmethod + def _upload_metadata_is_recent(cls, info: Dict[str, Any], cutoff_date: datetime) -> bool: + """Return True when upload metadata records activity inside retention.""" + for field in ("last_accessed", "created_at", "uploaded_at"): + parsed = cls._parse_upload_timestamp(info.get(field)) + if parsed is None: + continue + if parsed >= cutoff_date: + return True + return False + + @classmethod + def _upload_index_keys_for_file( + cls, + upload_index: Dict[str, Any], + upload_id: str, + file_path: str, + ) -> list[str]: + """Find a coherent set of index rows for one physical upload. + + Every related row must agree on ID, canonical path, owner, and a + non-empty checksum. Each row must also contain the complete lifecycle + timestamps written for new uploads. Ambiguous or incomplete index + state cannot authorize destructive cleanup. + """ + target_path = os.path.normcase(os.path.realpath(file_path)) + matches: list[str] = [] + owners: set[str] = set() + checksums: set[str] = set() + for key, info in upload_index.items(): + if not isinstance(info, dict): + continue + stored_path = info.get("path") + stored_real_path = ( + os.path.normcase(os.path.realpath(stored_path)) + if isinstance(stored_path, str) and stored_path + else None + ) + same_id = info.get("id") == upload_id + same_path = stored_real_path == target_path + if not same_id and not same_path: + continue + if not same_id or not same_path: + logger.warning( + "Skipping ambiguous cleanup candidate %s: related row has id=%r path=%r", + file_path, + info.get("id"), + stored_path, + ) + return [] + + owner = info.get("owner") + if not isinstance(owner, str) or not owner.strip(): + logger.warning( + "Skipping incomplete cleanup candidate %s: matching row has no owner", + file_path, + ) + return [] + + row_checksums = { + str(info.get(field)).strip().lower() + for field in ("hash", "checksum_sha256") + if info.get(field) is not None and str(info.get(field)).strip() + } + if not row_checksums: + logger.warning( + "Skipping incomplete cleanup candidate %s: matching row has no checksum", + file_path, + ) + return [] + if len(row_checksums) != 1: + logger.warning( + "Skipping ambiguous cleanup candidate %s: matching row has conflicting checksums", + file_path, + ) + return [] + + lifecycle_fields = ("uploaded_at", "created_at", "last_accessed") + if any( + cls._parse_upload_timestamp(info.get(field)) is None + for field in lifecycle_fields + ): + logger.warning( + "Skipping incomplete cleanup candidate %s: matching row lacks lifecycle timestamps", + file_path, + ) + return [] + + matches.append(key) + owners.add(owner) + checksums.update(row_checksums) + + if len(owners) > 1 or len(checksums) > 1: + logger.warning( + "Skipping ambiguous cleanup candidate %s: matching rows disagree on owner or checksum", + file_path, + ) + return [] + return matches + + def cleanup_old_uploads( + self, + referenced_upload_ids: Optional[set[str]] = None, + referenced_upload_hashes: Optional[set[str]] = None, + ): + """Remove expired uploads proven unreferenced by a complete snapshot. + + ``None`` means reference discovery was not completed, so cleanup fails + closed and removes nothing. The admin route supplies both sets after + scanning persisted chats, documents, and gallery records. + """ + if referenced_upload_ids is None or referenced_upload_hashes is None: + logger.warning("Upload cleanup skipped: persisted reference snapshot unavailable") + return 0 + + try: + cleanup_started_at = datetime.now() + cutoff_date = cleanup_started_at - timedelta(days=self.cleanup_days) cleaned_count = 0 - - for root, dirs, files in os.walk(self.upload_dir): - if root == self.upload_dir: - continue - - path_parts = root.split(os.sep) - if len(path_parts) >= 4: + + referenced_ids = {str(value) for value in referenced_upload_ids} + referenced_hashes = {str(value) for value in referenced_upload_hashes} + uploads_db_path = os.path.join(self.upload_dir, "uploads.json") + + # Keep index mutation and file removal serialized with upload writes. + # Each row removal is atomically persisted before the bytes are + # deleted; if deletion fails, the previous index is restored. + with self._index_lock: + current_index = dict(self._load_upload_index(fail_on_error=True)) + + for root, dirs, files in os.walk(self.upload_dir, followlinks=False): + is_junction = getattr(os.path, "isjunction", lambda _path: False) + dirs[:] = [ + directory + for directory in dirs + if not os.path.islink(os.path.join(root, directory)) + and not is_junction(os.path.join(root, directory)) + ] + if root == self.upload_dir: + continue + if not self._inside_upload_dir(root): + dirs[:] = [] + continue + + path_parts = root.split(os.sep) + if len(path_parts) < 4: + continue try: dir_date = datetime(int(path_parts[-3]), int(path_parts[-2]), int(path_parts[-1])) - if dir_date < cutoff_date: - for file in files: - file_path = os.path.join(root, file) - try: - os.remove(file_path) - cleaned_count += 1 - logger.info(f"Cleaned up old upload: {file_path}") - except Exception as e: - logger.warning(f"Failed to remove {file_path}: {e}") - - try: - os.rmdir(root) - logger.info(f"Removed empty upload directory: {root}") - except Exception as e: - logger.warning(f"Failed to remove directory {root}: {e}") except (ValueError, IndexError): continue - + if dir_date >= cutoff_date: + continue + + for file in files: + # Reference discovery only understands canonical upload + # IDs; unknown files fail closed instead of being swept. + if not self.validate_upload_id(file): + continue + + file_path = os.path.join(root, file) + if not self._inside_upload_dir(file_path): + logger.warning( + "Skipping cleanup candidate outside upload directory: %s", + file_path, + ) + continue + matching_keys = self._upload_index_keys_for_file( + current_index, + file, + file_path, + ) + matching_rows = [ + current_index[key] + for key in matching_keys + if isinstance(current_index.get(key), dict) + ] + + # Files without authoritative live index rows are not + # eligible for destructive cleanup. Reference hashes, + # recency, and ownership cannot be proven for them. + if not matching_rows: + continue + + is_referenced = file in referenced_ids or any( + str(info.get("id") or "") in referenced_ids + or str(info.get("hash") or "") in referenced_hashes + or str(info.get("checksum_sha256") or "") in referenced_hashes + for info in matching_rows + ) + metadata_is_recent = any( + self._upload_metadata_is_recent(info, cutoff_date) + for info in matching_rows + ) + if is_referenced or metadata_is_recent: + continue + + reduced_index = { + key: value + for key, value in current_index.items() + if key not in matching_keys + } + if matching_keys: + try: + self._atomic_write_json( + uploads_db_path, + reduced_index, + sync_backup=True, + ) + except Exception as e: + try: + self._atomic_write_json( + uploads_db_path, + current_index, + sync_backup=True, + ) + except Exception: + logger.exception( + "Failed to restore upload indexes after reconciliation failed for %s", + file_path, + ) + raise UploadCleanupSafetyError( + "upload index rollback failed before file removal" + ) from e + logger.warning( + "Failed to reconcile upload index before removing %s: %s", + file_path, + e, + ) + continue + + try: + os.remove(file_path) + except FileNotFoundError: + # The bytes are already absent. Keep the reduced + # lifecycle index instead of recreating a stale row. + current_index = reduced_index + logger.info( + "Reconciled missing expired upload from index: %s", + file_path, + ) + continue + except Exception as e: + if matching_keys: + try: + self._atomic_write_json( + uploads_db_path, + current_index, + sync_backup=True, + ) + except Exception: + logger.exception( + "Failed to restore upload index after removal failed for %s", + file_path, + ) + raise UploadCleanupSafetyError( + "upload index rollback failed after file removal was refused" + ) from e + logger.warning(f"Failed to remove {file_path}: {e}") + continue + + current_index = reduced_index + cleaned_count += 1 + logger.info(f"Cleaned up old unreferenced upload: {file_path}") + + try: + if not os.listdir(root): + os.rmdir(root) + logger.info(f"Removed empty upload directory: {root}") + except Exception as e: + logger.warning(f"Failed to inspect/remove directory {root}: {e}") + logger.info(f"Upload cleanup completed: {cleaned_count} files removed") return cleaned_count except Exception as e: logger.error(f"Upload cleanup failed: {e}") - return 0 + raise UploadCleanupSafetyError("upload cleanup safety checks failed") from e def validate_upload_id(self, upload_id: str) -> bool: """Validate that the upload ID matches the expected pattern.""" @@ -293,71 +672,103 @@ class UploadHandler: def _inside_upload_dir(self, path: str) -> bool: """Check if path is inside the upload directory.""" - base = os.path.realpath(self.upload_dir) - p = os.path.realpath(path) + base = os.path.normcase(os.path.realpath(self.upload_dir)) + p = os.path.normcase(os.path.realpath(path)) try: return os.path.commonpath([base, p]) == base except Exception: return False - def _atomic_write_json(self, path: str, data: dict) -> None: + def _atomic_write_json( + self, + path: str, + data: dict, + *, + sync_backup: bool = False, + ) -> None: """Write `data` to `path` atomically: write to a temp file in the same directory, then `os.replace` onto the target. The kernel guarantees `os.replace` is atomic on POSIX, so a reader either sees the old contents or the new contents, never a half-written - file. Also keeps a `.bak` sibling of the previous good state. + file. Normally `.bak` retains the previous good state. Destructive + lifecycle transitions use ``sync_backup=True`` so recovery cannot + resurrect metadata for bytes that were deliberately removed. """ directory = os.path.dirname(path) or "." - fd, tmp = tempfile.mkstemp(prefix=".uploads-", suffix=".tmp", dir=directory) - try: - with os.fdopen(fd, "w", encoding="utf-8") as f: - json.dump(data, f, indent=2) - f.flush() - os.fsync(f.fileno()) - if os.path.exists(path): - bak = path + ".bak" + + def _replace_json(target: str) -> None: + fd, tmp = tempfile.mkstemp( + prefix=".uploads-", + suffix=".tmp", + dir=directory, + ) + try: + with os.fdopen(fd, "w", encoding="utf-8") as f: + json.dump(data, f, indent=2) + f.flush() + os.fsync(f.fileno()) + os.replace(tmp, target) + except Exception: try: - shutil.copy2(path, bak) + os.unlink(tmp) except OSError: pass - os.replace(tmp, path) - # Update cache if this is the main index - if path.endswith("uploads.json"): - self._index_cache = data - try: - self._index_mtime = os.path.getmtime(path) - except OSError: - self._index_mtime = time.time() - except Exception: + raise + + if sync_backup: + _replace_json(path + ".bak") + elif os.path.exists(path): try: - os.unlink(tmp) + shutil.copy2(path, path + ".bak") except OSError: pass - raise - def _load_upload_index(self) -> Dict[str, Any]: + _replace_json(path) + # Update cache if this is the main index + if path.endswith("uploads.json"): + self._index_cache = data + try: + self._index_mtime = os.path.getmtime(path) + except OSError: + self._index_mtime = time.time() + + def _load_upload_index(self, *, fail_on_error: bool = False) -> Dict[str, Any]: """Load the upload index from disk/cache. Uses mtime-based validation - to avoid redundant parsing on hot paths. + to avoid redundant parsing on hot paths. When ``fail_on_error`` is + true, a missing, malformed, or unreadable live index raises so + destructive callers cannot mistake corruption for an empty store. """ uploads_db_path = os.path.join(self.upload_dir, "uploads.json") - if not os.path.exists(uploads_db_path): + candidates = (uploads_db_path, uploads_db_path + ".bak") + if fail_on_error: + # A backup is intentionally the previous snapshot. It is useful for + # non-destructive reads, but cannot authorize deletion when the live + # index is missing or corrupt. + if not os.path.exists(uploads_db_path): + raise ValueError("live uploads database is missing") + existing_candidates = [uploads_db_path] + else: + existing_candidates = [path for path in candidates if os.path.exists(path)] + if not existing_candidates: self._index_cache = {} self._index_mtime = 0.0 return {} # Check cache validity try: - mtime = os.path.getmtime(uploads_db_path) - if self._index_cache is not None and mtime <= self._index_mtime: + mtime = max(os.path.getmtime(path) for path in existing_candidates) + if ( + not fail_on_error + and self._index_cache is not None + and mtime <= self._index_mtime + ): return self._index_cache except OSError: mtime = 0.0 # Try the live file first, fall back to the .bak sibling if the # live file is truncated/corrupted. - for candidate in (uploads_db_path, uploads_db_path + ".bak"): - if not os.path.exists(candidate): - continue + for candidate in existing_candidates: try: with open(candidate, "r", encoding="utf-8") as f: data = json.load(f) @@ -369,6 +780,8 @@ class UploadHandler: logger.warning(f"Failed to read uploads database ({candidate}): {e}") continue + if fail_on_error: + raise ValueError("live uploads database is unreadable") self._index_cache = {} return {} @@ -381,6 +794,144 @@ class UploadHandler: return dict(info) return None + def reserve_upload( + self, + upload_id: str, + *, + owner: Optional[str], + auth_manager: Any = None, + allow_admin: bool = False, + ) -> Optional[Dict[str, Any]]: + """Owner-check and reserve an indexed upload against cleanup. + + The live index lookup, ownership/path validation, and access touch all + occur under the cleanup lock. A durable-reference writer must not + commit when this returns ``None``. + """ + if not self.validate_upload_id(upload_id): + return None + + auth_configured = bool(auth_manager and getattr(auth_manager, "is_configured", False)) + if auth_configured and not owner: + return None + + uploads_db_path = os.path.join(self.upload_dir, "uploads.json") + with self._index_lock: + try: + current = dict(self._load_upload_index(fail_on_error=True)) + except Exception: + logger.warning("Cannot reserve upload %s without a valid live index", upload_id) + return None + matching_keys = [ + key + for key, info in current.items() + if isinstance(info, dict) and info.get("id") == upload_id + ] + if not matching_keys: + return None + + matching_rows = [dict(current[key]) for key in matching_keys] + row_owners = { + str(row.get("owner")) if row.get("owner") is not None else None + for row in matching_rows + } + row_hashes = { + str(row.get("hash") or row.get("checksum_sha256")) + for row in matching_rows + if row.get("hash") or row.get("checksum_sha256") + } + if len(row_owners) != 1 or len(row_hashes) > 1: + logger.warning( + "Cannot reserve ambiguous upload index rows for %s", + upload_id, + ) + return None + + is_admin = False + if allow_admin and owner and auth_manager and hasattr(auth_manager, "is_admin"): + try: + is_admin = bool(auth_manager.is_admin(owner)) + except Exception: + is_admin = False + + now = datetime.now() + current_info = matching_rows[0] + if owner and not is_admin and current_info.get("owner") != owner: + return None + if not owner and current_info.get("owner") is not None: + return None + + existing_paths: set[str] = set() + for row in matching_rows: + stored_path = row.get("path") + if not stored_path: + continue + if not self._inside_upload_dir(stored_path): + logger.warning( + "Cannot reserve upload %s with an out-of-root index path", + upload_id, + ) + return None + if os.path.isfile(stored_path): + if os.path.basename(stored_path) != upload_id: + return None + existing_paths.add(os.path.normcase(os.path.realpath(stored_path))) + if len(existing_paths) > 1: + logger.warning("Cannot reserve upload %s with multiple indexed paths", upload_id) + return None + path = next(iter(existing_paths), None) or self._find_upload_path(upload_id) + if not path or not os.path.isfile(path) or not self._inside_upload_dir(path): + return None + + last_accessed = self._parse_upload_timestamp(current_info.get("last_accessed")) + path_changed = current_info.get("path") != path + needs_write = ( + path_changed + or last_accessed is None + or last_accessed < now - timedelta(minutes=5) + ) + if needs_write: + accessed_at = now.isoformat() + updated_index = dict(current) + for key in matching_keys: + updated = dict(updated_index[key]) + updated["path"] = path + updated["last_accessed"] = accessed_at + updated_index[key] = updated + try: + self._atomic_write_json( + uploads_db_path, + updated_index, + sync_backup=True, + ) + except Exception: + try: + self._atomic_write_json( + uploads_db_path, + current, + sync_backup=True, + ) + except Exception: + logger.exception( + "Failed to restore upload indexes after reservation failed for %s", + upload_id, + ) + logger.exception("Failed to reserve upload %s against cleanup", upload_id) + return None + current_info = dict(updated_index[matching_keys[0]]) + + resolved = dict(current_info) + resolved.setdefault("id", upload_id) + resolved["path"] = path + resolved.setdefault("name", os.path.basename(path)) + resolved.setdefault("original_name", resolved["name"]) + resolved.setdefault("mime", mimetypes.guess_type(path)[0] or "application/octet-stream") + if resolved.get("hash") and not resolved.get("checksum_sha256"): + resolved["checksum_sha256"] = resolved["hash"] + if resolved.get("uploaded_at") and not resolved.get("created_at"): + resolved["created_at"] = resolved["uploaded_at"] + return resolved + def _renamed_upload_index_key(self, key: str, info: Dict[str, Any], old_owner: str, new_owner: str) -> str: """Return the storage key to use after renaming an owned upload row. @@ -475,16 +1026,32 @@ class UploadHandler: if not self.validate_upload_id(upload_id): return None + candidates: list[str] = [] direct = os.path.join(self.upload_dir, upload_id) - if os.path.exists(direct) and self._inside_upload_dir(direct): - return direct + if os.path.isfile(direct) and self._inside_upload_dir(direct): + candidates.append(os.path.realpath(direct)) - for root, _dirs, files in os.walk(self.upload_dir, followlinks=False): + for root, dirs, files in os.walk(self.upload_dir, followlinks=False): + is_junction = getattr(os.path, "isjunction", lambda _path: False) + dirs[:] = [ + directory + for directory in dirs + if not os.path.islink(os.path.join(root, directory)) + and not is_junction(os.path.join(root, directory)) + ] if upload_id in files: path = os.path.join(root, upload_id) - if self._inside_upload_dir(path): - return path - return None + if os.path.isfile(path) and self._inside_upload_dir(path): + real_path = os.path.realpath(path) + if real_path not in candidates: + candidates.append(real_path) + if len(candidates) > 1: + logger.warning( + "Upload ID %s resolves to multiple physical files", + upload_id, + ) + return None + return candidates[0] if candidates else None def resolve_upload( self, @@ -493,52 +1060,20 @@ class UploadHandler: auth_manager: Any = None, allow_admin: bool = True, ) -> Optional[Dict[str, Any]]: - """Resolve an upload ID to metadata only if the caller may read it. + """Resolve and reserve an upload only if the caller may read it. This is the owner-aware lookup used by internal processors. Public download routes already perform owner checks; chat/document paths must - do the same before reading file bytes server-side. + do the same before reading file bytes server-side. Reservation shares + cleanup's lifecycle lock and prevents a newly persisted reference from + racing final deletion. """ - if not self.validate_upload_id(upload_id): - logger.warning(f"Invalid upload ID format: {upload_id}") - return None - - auth_configured = bool(auth_manager and getattr(auth_manager, "is_configured", False)) - if auth_configured and not owner: - return None - - info = self.get_upload_info(upload_id) or {} - is_admin = False - if allow_admin and owner and auth_manager and hasattr(auth_manager, "is_admin"): - try: - is_admin = bool(auth_manager.is_admin(owner)) - except Exception: - is_admin = False - - if owner and not is_admin: - if info.get("owner") != owner: - logger.warning("Upload %s denied for owner %s", upload_id, owner) - return None - if not owner and info.get("owner") is not None: - logger.warning("Upload %s denied without an authenticated owner", upload_id) - return None - - path = info.get("path") - if not path or not os.path.exists(path) or not self._inside_upload_dir(path): - path = self._find_upload_path(upload_id) - if not path: - return None - if not self._inside_upload_dir(path): - logger.warning(f"Upload path outside upload directory: {path}") - return None - - resolved = dict(info) - resolved.setdefault("id", upload_id) - resolved["path"] = path - resolved.setdefault("name", os.path.basename(path)) - resolved.setdefault("original_name", resolved["name"]) - resolved.setdefault("mime", mimetypes.guess_type(path)[0] or "application/octet-stream") - return resolved + return self.reserve_upload( + upload_id, + owner=owner, + auth_manager=auth_manager, + allow_admin=allow_admin, + ) def cleanup_rate_limits(self): """Remove stale entries from upload_rate_log.""" @@ -706,6 +1241,9 @@ class UploadHandler: # fresh-insert path below; release the lock first. raise LookupError("upload entry vanished mid-dedupe") existing_file["last_accessed"] = datetime.now().isoformat() + existing_file.setdefault("checksum_sha256", file_hash) + if existing_file.get("uploaded_at"): + existing_file.setdefault("created_at", existing_file["uploaded_at"]) current[live_key] = existing_file self._atomic_write_json(uploads_db_path, current) except LookupError: @@ -721,7 +1259,9 @@ class UploadHandler: "size": existing_file["size"], "name": existing_file["original_name"], "hash": file_hash, + "checksum_sha256": existing_file.get("checksum_sha256") or file_hash, "uploaded_at": existing_file["uploaded_at"], + "created_at": existing_file.get("created_at") or existing_file["uploaded_at"], "owner": existing_file.get("owner"), "width": existing_file.get("width"), "height": existing_file.get("height"), @@ -744,6 +1284,7 @@ class UploadHandler: raise HTTPException(status_code=500, detail=f"Failed to save file: {str(e)}") # Create file metadata + created_at = datetime.now().isoformat() file_metadata = { "id": file_id, "path": file_path, @@ -751,9 +1292,11 @@ class UploadHandler: "size": file_size, "name": safe_filename, "hash": file_hash, + "checksum_sha256": file_hash, "original_name": original_filename, - "uploaded_at": datetime.now().isoformat(), - "last_accessed": datetime.now().isoformat(), + "uploaded_at": created_at, + "created_at": created_at, + "last_accessed": created_at, "client_ip": client_ip, "owner": owner, } diff --git a/static/js/chat.js b/static/js/chat.js index f9a035a8c..06c7a1dc7 100644 --- a/static/js/chat.js +++ b/static/js/chat.js @@ -1822,7 +1822,7 @@ import { wireArrowUpRecall, getLastUserMessageFromChatHistory } from './composer typewriterInto(roundHolder.querySelector('.body'), errMsg); break; } - if (json.delta || json.type === 'agent_prep' || json.type === 'tool_start' || json.type === 'tool_output' || json.type === 'tool_progress' || json.type === 'agent_step' || json.type === 'doc_stream_open' || json.type === 'doc_stream_delta' || json.type === 'research_progress') { + if (json.delta || json.type === 'agent_prep' || json.type === 'tool_start' || json.type === 'tool_output' || json.type === 'tool_progress' || json.type === 'agent_step' || json.type === 'loop_breaker_triggered' || json.type === 'intent_nudge_exhausted' || json.type === 'doc_stream_open' || json.type === 'doc_stream_delta' || json.type === 'research_progress') { clearResponseTimeout(); clearProcessingProbe(); clearFirstTokenWaitTimers(); @@ -2852,6 +2852,22 @@ import { wireArrowUpRecall, getLastUserMessageFromChatHistory } from './composer const chatBox = document.getElementById('chat-history'); chatBox.appendChild(budgetDiv); + } else if (json.type === 'loop_breaker_triggered' || json.type === 'intent_nudge_exhausted') { + if (_isBg) continue; + _cancelThinkingTimer(); + _removeThinkingSpinner(); + const guardDiv = document.createElement('div'); + guardDiv.className = 'stopped-indicator'; + const guardLabel = document.createElement('span'); + guardLabel.textContent = `[Agent guard: ${json.message || json.reason || 'internal stop'}]`; + guardDiv.appendChild(guardLabel); + const targetBody = roundHolder && roundHolder.querySelector('.body'); + if (targetBody) targetBody.appendChild(guardDiv); + else { + const chatBox = document.getElementById('chat-history'); + if (chatBox) chatBox.appendChild(guardDiv); + } + } else if (json.type === 'teacher_takeover') { if (_isBg) continue; _cancelThinkingTimer(); diff --git a/tests/test_agent_rounds_exhausted.py b/tests/test_agent_rounds_exhausted.py index 178faa8c1..b79245170 100644 --- a/tests/test_agent_rounds_exhausted.py +++ b/tests/test_agent_rounds_exhausted.py @@ -68,3 +68,24 @@ def test_no_rounds_exhausted_on_normal_finish(monkeypatch): # A plain answer (no tool block) -> done-break on round 1 -> no event. events = _run_loop(monkeypatch, "All done, here is your answer.", max_rounds=2) assert not any(e.get("type") == "rounds_exhausted" for e in events), events + + +def test_emits_intent_nudge_exhausted_when_cap_is_exhausted(monkeypatch): + _patch_common(monkeypatch) + + events = _run_loop(monkeypatch, "Let me check the logs", max_rounds=5) + + guard = next((e for e in events if e.get("type") == "intent_nudge_exhausted"), None) + assert guard is not None, events + assert guard["reason"] == "intent_without_action_nudge_cap" + assert guard["nudges"] == 2 + + +def test_emits_loop_breaker_triggered_when_loop_breaker_trips(monkeypatch): + _patch_common(monkeypatch) + + events = _run_loop(monkeypatch, "```bash\necho hi\n```", max_rounds=6) + + guard = next((e for e in events if e.get("type") == "loop_breaker_triggered"), None) + assert guard is not None, events + assert guard["reason"] == "loop_breaker_stall" diff --git a/tests/test_attachment_refs.py b/tests/test_attachment_refs.py new file mode 100644 index 000000000..0efd161c5 --- /dev/null +++ b/tests/test_attachment_refs.py @@ -0,0 +1,75 @@ +import json + +from src.attachment_refs import ( + attachment_ref, + persistable_message_content, + search_index_text, +) + + +def test_persistable_message_content_replaces_inline_media_with_attachment_ref(): + metadata = { + "attachments": [ + { + "id": "abc123.png", + "name": "diagram.png", + "mime": "image/png", + "size": 42, + "checksum_sha256": "sha256-digest", + "created_at": "2026-07-09T12:00:00", + "vision": "A small architecture diagram.", + } + ] + } + content = [ + {"type": "text", "text": "Please inspect this."}, + { + "type": "image_url", + "image_url": {"url": "data:image/png;base64," + ("A" * 5000)}, + }, + ] + + stored = persistable_message_content(content, metadata) + + assert "base64" not in stored + assert "A" * 100 not in stored + assert "Please inspect this." in stored + assert "Attachment: diagram.png" in stored + assert "id=abc123.png" in stored + assert "sha256=sha256-digest" in stored + assert "A small architecture diagram." in stored + + +def test_search_index_text_strips_legacy_serialized_data_url_blocks(): + legacy = json.dumps([ + {"type": "text", "text": "Find this useful caption"}, + { + "type": "image_url", + "image_url": {"url": "data:image/jpeg;base64," + ("B" * 4096)}, + }, + ]) + + indexed = search_index_text(legacy) + + assert indexed == "Find this useful caption\n[1 inline media payload omitted]" + + +def test_attachment_ref_normalizes_hash_aliases(): + ref = attachment_ref({ + "id": "file-id", + "original_name": "report.pdf", + "mime": "application/pdf", + "size": 99, + "hash": "abc", + "uploaded_at": "2026-07-09T12:00:00", + }) + + assert ref == { + "type": "attachment_ref", + "attachment_id": "file-id", + "name": "report.pdf", + "mime": "application/pdf", + "size": 99, + "checksum_sha256": "abc", + "created_at": "2026-07-09T12:00:00", + } diff --git a/tests/test_chat_helpers.py b/tests/test_chat_helpers.py index 0e2cce1f7..d2d5b2673 100644 --- a/tests/test_chat_helpers.py +++ b/tests/test_chat_helpers.py @@ -233,6 +233,10 @@ def test_build_uploaded_file_manifest_filters_and_nulls_unreadable_paths(monkeyp ) assert [item["id"] for item in manifest] == ["good", "outside", "missing"] + assert manifest[0]["type"] == "attachment_ref" + assert manifest[0]["attachment_id"] == "good" + assert manifest[0]["uri"] == "odysseus://attachment/good" + assert manifest[0]["read_policy"] == "owner_checked_upload" assert os.path.realpath(manifest[0]["path"]) == os.path.realpath(good) assert manifest[1]["path"] is None assert manifest[2]["path"] is None diff --git a/tests/test_parse_msg_content_jsonlike_string.py b/tests/test_parse_msg_content_jsonlike_string.py index 87d44fe2b..76a7d5ece 100644 --- a/tests/test_parse_msg_content_jsonlike_string.py +++ b/tests/test_parse_msg_content_jsonlike_string.py @@ -1,13 +1,9 @@ -"""A plain text message that merely *looks* like a JSON array of objects must -NOT be silently re-parsed into a list on reload. +"""Persistence contracts for JSON-like text and multimodal chat content. -_parse_msg_content de-serializes multimodal (image/audio) content back into a -list of content blocks. The old heuristic accepted ANY string that started -with "[{" and contained the substring '"type"'. A user who pasted an API -schema / sample such as `[{"type": "object", "name": "foo"}]` therefore had -their text message permanently corrupted into a Python list on the next -session hydration. The fix restricts the round-trip to lists whose elements -are all recognized content-block types (text/image_url/audio/...). +Plain text that resembles a JSON content-block list must remain an exact +string. Real provider multimodal blocks follow the durable attachment +contract: readable text plus stable attachment metadata is persisted, while +raw inline media bytes are omitted. """ import tempfile import uuid @@ -65,16 +61,48 @@ def test_jsonlike_user_string_not_corrupted(manager): assert reloaded.history[0].content == text -def test_real_multimodal_content_still_round_trips(manager): +def test_real_multimodal_content_persists_reference_without_base64(manager): sid = "sess-" + uuid.uuid4().hex[:8] _make_session(sid) + attachment_id = "a" * 32 + ".png" multimodal = [ {"type": "text", "text": "what is this?"}, {"type": "image_url", "image_url": {"url": "data:image/png;base64,AAAA"}}, ] - msgs = [ChatMessage(role="user", content=multimodal)] + metadata = { + "attachments": [ + { + "id": attachment_id, + "name": "diagram.png", + "mime": "image/png", + "size": 4, + "checksum_sha256": "sha256-digest", + } + ] + } + msgs = [ChatMessage(role="user", content=multimodal, metadata=metadata)] assert manager.replace_messages(sid, msgs) is True + expected = ( + "what is this?\n" + "[1 inline media payload omitted]\n" + f"[Attachment: diagram.png | id={attachment_id} | mime=image/png | " + "size=4 bytes | sha256=sha256-digest]" + ) + + db = _TS() + try: + stored = db.query(cdb.ChatMessage).filter_by(session_id=sid).one() + assert stored.content == expected + assert "what is this?" in stored.content + assert attachment_id in stored.content + assert "data:image/png;base64,AAAA" not in stored.content + assert "base64" not in stored.content + assert "AAAA" not in stored.content + finally: + db.close() + manager.sessions.clear() reloaded = manager.get_session(sid) - assert reloaded.history[0].content == multimodal + assert reloaded.history[0].content == expected + assert reloaded.history[0].metadata["attachments"][0]["id"] == attachment_id diff --git a/tests/test_replace_messages_multimodal.py b/tests/test_replace_messages_multimodal.py index ec8951577..19a668fa7 100644 --- a/tests/test_replace_messages_multimodal.py +++ b/tests/test_replace_messages_multimodal.py @@ -1,14 +1,9 @@ -"""replace_messages must JSON-serialize multimodal (list) content. +"""replace_messages must persist readable, path-free multimodal history. -A chat with an image/audio attachment carries list content. When such a -chat is compacted, the manual-compaction path calls replace_messages with -the retained messages. replace_messages wrote message.content straight into -the Text column, so SQLAlchemy bound the list\'s single-quoted repr. On -reload _parse_msg_content only de-serializes a string that contains the -double-quoted "type", so the repr failed the check and the message came -back as a corrupted string blob - the attachment was destroyed. The -sibling _persist_message json.dumps-es list content; replace_messages did -not. +Live model input may contain provider-specific media blocks and inline data +URLs. Compaction uses replace_messages for the retained transcript, which must +store readable text plus stable structured attachment references without +copying raw base64 payloads into ChatMessage.content. """ import uuid @@ -27,6 +22,7 @@ def manager(monkeypatch): monkeypatch.setattr(sm, "SessionLocal", _TS) mgr = sm.SessionManager.__new__(sm.SessionManager) mgr.sessions = {} + mgr.upload_handler = None return mgr @@ -41,33 +37,71 @@ def _make_session(sid, owner="alice"): db.close() -def test_multimodal_content_round_trips_through_replace_messages(manager): +def test_multimodal_content_persists_text_and_attachment_ref_without_payload(manager): sid = "sess-" + uuid.uuid4().hex[:8] _make_session(sid) + upload_id = "a" * 32 + ".png" multimodal = [ {"type": "text", "text": "what is this?"}, {"type": "image_url", "image_url": {"url": "data:image/png;base64,AAAA"}}, ] - msgs = [ChatMessage(role="user", content=multimodal)] + msgs = [ChatMessage( + role="user", + content=multimodal, + metadata={ + "attachments": [{ + "id": upload_id, + "name": "diagram.png", + "mime": "image/png", + "size": 4, + "checksum_sha256": "sha256-digest", + }] + }, + )] assert manager.replace_messages(sid, msgs) is True + expected = ( + "what is this?\n" + "[1 inline media payload omitted]\n" + f"[Attachment: diagram.png | id={upload_id} | mime=image/png | " + "size=4 bytes | sha256=sha256-digest]" + ) + + db = _TS() + try: + stored = db.query(cdb.ChatMessage).filter_by(session_id=sid).one() + assert stored.content == expected + assert "data:image/png;base64,AAAA" not in stored.content + assert "base64" not in stored.content + assert "AAAA" not in stored.content + finally: + db.close() + # Drop the in-memory cache so the next read hydrates from the DB. manager.sessions.clear() reloaded = manager.get_session(sid) assert len(reloaded.history) == 1 - # Content must come back as the original list, not a repr string blob. - assert reloaded.history[0].content == multimodal + persisted = reloaded.history[0].content + assert isinstance(persisted, str) + assert persisted == expected + assert reloaded.history[0].metadata["attachments"][0]["id"] == upload_id + assert ( + reloaded.history[0].metadata["attachments"][0]["checksum_sha256"] + == "sha256-digest" + ) -def test_plain_string_content_still_round_trips(manager): +def test_jsonlike_plain_string_content_still_round_trips(manager): sid = "sess-" + uuid.uuid4().hex[:8] _make_session(sid) - msgs = [ChatMessage(role="user", content="just text")] + text = '[{"type": "object", "name": "foo"}]' + msgs = [ChatMessage(role="user", content=text)] assert manager.replace_messages(sid, msgs) is True manager.sessions.clear() reloaded = manager.get_session(sid) - assert reloaded.history[0].content == "just text" + assert isinstance(reloaded.history[0].content, str) + assert reloaded.history[0].content == text def test_replace_messages_keeps_history_alias_for_context_messages(manager): diff --git a/tests/test_replace_messages_upload_reservations.py b/tests/test_replace_messages_upload_reservations.py new file mode 100644 index 000000000..37a86d16e --- /dev/null +++ b/tests/test_replace_messages_upload_reservations.py @@ -0,0 +1,259 @@ +"""Upload lifecycle guarantees for compaction's replace_messages path.""" + +import concurrent.futures +import json +import os +import threading +import uuid + +import pytest +from sqlalchemy import event + +import core.database as cdb +import core.session_manager as session_manager_module +from core.models import ChatMessage +from src.upload_handler import UploadHandler +from tests.helpers.sqlite_db import make_temp_sqlite + + +OLD_TIMESTAMP = "2000-01-01T00:00:00" + + +@pytest.fixture +def manager_db(monkeypatch): + SessionLocal, engine, tmpfile = make_temp_sqlite(cdb.Base.metadata) + monkeypatch.setattr(session_manager_module, "SessionLocal", SessionLocal) + manager = session_manager_module.SessionManager.__new__( + session_manager_module.SessionManager + ) + manager.sessions = {} + manager.upload_handler = None + try: + yield manager, SessionLocal, engine + finally: + engine.dispose() + tmpfile.close() + try: + os.unlink(tmpfile.name) + except OSError: + pass + + +def _seed_session(SessionLocal, *, owner="alice", content="existing durable history"): + session_id = "replace-" + uuid.uuid4().hex + db = SessionLocal() + try: + db.add(cdb.Session( + id=session_id, + owner=owner, + name="Compaction reservation regression", + model="test-model", + endpoint_url="http://localhost:11434", + archived=False, + message_count=1, + )) + db.add(cdb.ChatMessage( + id="message-" + uuid.uuid4().hex, + session_id=session_id, + role="user", + content=content, + meta_data=json.dumps({"source": "before-replacement"}), + )) + db.commit() + finally: + db.close() + return session_id + + +def _attachment_message(upload_id, text): + return ChatMessage( + role="user", + content=text, + metadata={ + "attachments": [{ + "id": upload_id, + "name": f"{text}.txt", + "mime": "text/plain", + "size": len(text), + }] + }, + ) + + +def _durable_messages(SessionLocal, session_id): + db = SessionLocal() + try: + return [ + (message.role, message.content, message.meta_data) + for message in db.query(cdb.ChatMessage) + .filter(cdb.ChatMessage.session_id == session_id) + .order_by(cdb.ChatMessage.timestamp, cdb.ChatMessage.id) + .all() + ] + finally: + db.close() + + +def test_replace_messages_reserves_every_incoming_attachment_before_delete( + manager_db, + monkeypatch, +): + manager, SessionLocal, engine = manager_db + session_id = _seed_session(SessionLocal) + manager.upload_handler = object() + incoming = [ + _attachment_message("1" * 32 + ".txt", "first"), + _attachment_message("2" * 32 + ".txt", "second"), + ] + message_mutations = [] + reservation_calls = [] + + def record_sql(_conn, _cursor, statement, _parameters, _context, _executemany): + normalized = statement.lstrip().upper() + if normalized.startswith(("DELETE FROM CHAT_MESSAGES", "INSERT INTO CHAT_MESSAGES")): + message_mutations.append(normalized.split(maxsplit=1)[0]) + + def reserve(handler, owner, content, metadata): + assert message_mutations == [] + reservation_calls.append((handler, owner, content, metadata)) + return None + + event.listen(engine, "before_cursor_execute", record_sql) + monkeypatch.setattr( + session_manager_module, + "reserve_message_upload_references", + reserve, + ) + try: + assert manager.replace_messages(session_id, incoming) is True + finally: + event.remove(engine, "before_cursor_execute", record_sql) + + assert [call[2] for call in reservation_calls] == ["first", "second"] + assert all(call[0] is manager.upload_handler for call in reservation_calls) + assert all(call[1] == "alice" for call in reservation_calls) + assert message_mutations == ["DELETE", "INSERT"] + + +def test_replace_messages_reservation_failure_leaves_durable_history_unchanged( + manager_db, + monkeypatch, +): + manager, SessionLocal, engine = manager_db + session_id = _seed_session(SessionLocal) + before = _durable_messages(SessionLocal, session_id) + manager.upload_handler = object() + missing_upload_id = "4" * 32 + ".txt" + incoming = [ + _attachment_message("3" * 32 + ".txt", "available"), + _attachment_message(missing_upload_id, "missing"), + ] + calls = [] + message_mutations = [] + + def record_sql(_conn, _cursor, statement, _parameters, _context, _executemany): + normalized = statement.lstrip().upper() + if normalized.startswith(("DELETE FROM CHAT_MESSAGES", "INSERT INTO CHAT_MESSAGES")): + message_mutations.append(normalized.split(maxsplit=1)[0]) + + def reserve(_handler, _owner, content, _metadata): + calls.append(content) + return missing_upload_id if content == "missing" else None + + event.listen(engine, "before_cursor_execute", record_sql) + monkeypatch.setattr( + session_manager_module, + "reserve_message_upload_references", + reserve, + ) + try: + assert manager.replace_messages(session_id, incoming) is False + finally: + event.remove(engine, "before_cursor_execute", record_sql) + + assert calls == ["available", "missing"] + assert message_mutations == [] + assert _durable_messages(SessionLocal, session_id) == before + assert [message.content for message in manager.sessions[session_id].history] == [ + "existing durable history" + ] + assert all("_db_id" not in (message.metadata or {}) for message in incoming) + + +def test_cleanup_cannot_delete_attachment_during_concurrent_compaction_replacement( + manager_db, + monkeypatch, + tmp_path, +): + manager, SessionLocal, _engine = manager_db + session_id = _seed_session(SessionLocal) + base_dir = tmp_path / "base" + upload_dir = tmp_path / "uploads" + base_dir.mkdir() + upload_dir.mkdir() + handler = UploadHandler(str(base_dir), str(upload_dir)) + manager.upload_handler = handler + + upload_id = "5" * 32 + ".txt" + upload_hash = "6" * 64 + dated_dir = upload_dir / "2000" / "01" / "01" + dated_dir.mkdir(parents=True) + upload_path = dated_dir / upload_id + upload_path.write_text("attachment retained by compaction", encoding="utf-8") + upload_row = { + "id": upload_id, + "path": str(upload_path), + "mime": "text/plain", + "size": upload_path.stat().st_size, + "name": "compaction.txt", + "original_name": "compaction.txt", + "hash": upload_hash, + "checksum_sha256": upload_hash, + "uploaded_at": OLD_TIMESTAMP, + "created_at": OLD_TIMESTAMP, + "last_accessed": OLD_TIMESTAMP, + "owner": "alice", + } + (upload_dir / "uploads.json").write_text( + json.dumps({f"alice:{upload_hash}": upload_row}), + encoding="utf-8", + ) + handler._index_cache = None + + reservation_write_entered = threading.Event() + release_reservation_write = threading.Event() + real_atomic_write = handler._atomic_write_json + + def block_reservation_write(path, data, *, sync_backup=False): + refreshed = any( + isinstance(row, dict) + and row.get("id") == upload_id + and row.get("last_accessed") != OLD_TIMESTAMP + for row in data.values() + ) + if sync_backup and refreshed and not reservation_write_entered.is_set(): + reservation_write_entered.set() + assert release_reservation_write.wait(5) + return real_atomic_write(path, data, sync_backup=sync_backup) + + monkeypatch.setattr(handler, "_atomic_write_json", block_reservation_write) + incoming = [_attachment_message(upload_id, "retained after compaction")] + + with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool: + replace_future = pool.submit(manager.replace_messages, session_id, incoming) + assert reservation_write_entered.wait(5) + cleanup_future = pool.submit(handler.cleanup_old_uploads, set(), set()) + try: + with pytest.raises(concurrent.futures.TimeoutError): + cleanup_future.result(timeout=0.1) + finally: + release_reservation_write.set() + + assert replace_future.result(timeout=5) is True + assert cleanup_future.result(timeout=5) == 0 + + assert upload_path.is_file() + assert handler.resolve_upload(upload_id, owner="alice") is not None + durable = _durable_messages(SessionLocal, session_id) + assert len(durable) == 1 + assert json.loads(durable[0][2])["attachments"][0]["id"] == upload_id diff --git a/tests/test_upload_handler_cleanup.py b/tests/test_upload_handler_cleanup.py new file mode 100644 index 000000000..9810c2b20 --- /dev/null +++ b/tests/test_upload_handler_cleanup.py @@ -0,0 +1,831 @@ +import asyncio +import concurrent.futures +import json +import os +import threading +from datetime import datetime +from pathlib import Path +from types import SimpleNamespace + +import pytest +from fastapi import HTTPException + +from core.database import ( + Base, + ChatMessage as DbChatMessage, + CalendarCal, + CalendarEvent, + Document, + DocumentVersion, + GalleryImage, + Note, + Session as DbSession, +) +from src.upload_handler import ( + UploadCleanupSafetyError, + UploadHandler, + extract_internal_upload_ids, + reserve_message_upload_references, + reserve_upload_references, +) +from tests.helpers.sqlite_db import make_temp_sqlite + + +OLD_TIMESTAMP = "2000-01-01T00:00:00" + + +class _AdminAuth: + is_configured = True + + @staticmethod + def is_admin(user): + return user == "admin" + + +class _AdminRequest: + headers = {} + state = SimpleNamespace(current_user="admin") + app = SimpleNamespace(state=SimpleNamespace(auth_manager=_AdminAuth())) + + +def _make_handler(tmp_path: Path) -> UploadHandler: + base_dir = tmp_path / "base" + upload_dir = tmp_path / "uploads" + base_dir.mkdir() + upload_dir.mkdir() + return UploadHandler(str(base_dir), str(upload_dir)) + + +def _seed_old_uploads(handler: UploadHandler, rows: list[dict]) -> dict[str, Path]: + dated_dir = Path(handler.upload_dir) / "2000" / "01" / "01" + dated_dir.mkdir(parents=True) + index = {} + paths = {} + for row in rows: + upload_id = row["id"] + path = dated_dir / upload_id + path.write_bytes(row.get("bytes", upload_id.encode("ascii"))) + info = { + "id": upload_id, + "path": str(path), + "mime": row.get("mime", "application/octet-stream"), + "size": path.stat().st_size, + "name": row.get("name", upload_id), + "original_name": row.get("name", upload_id), + "hash": row["hash"], + "checksum_sha256": row["hash"], + "uploaded_at": row.get("uploaded_at", OLD_TIMESTAMP), + "created_at": row.get("created_at", OLD_TIMESTAMP), + "last_accessed": row.get("last_accessed", OLD_TIMESTAMP), + "owner": row.get("owner", "alice"), + } + index[f"{info['owner']}:{info['hash']}"] = info + paths[upload_id] = path + + Path(handler.upload_dir, "uploads.json").write_text( + json.dumps(index), + encoding="utf-8", + ) + handler._index_cache = None + return paths + + +def _manual_cleanup_endpoint(handler: UploadHandler, monkeypatch): + import fastapi.dependencies.utils as dependency_utils + from routes.upload_routes import router, setup_upload_routes + + monkeypatch.setattr(dependency_utils, "ensure_multipart_is_installed", lambda: None) + before = len(router.routes) + setup_upload_routes(handler) + return { + route.endpoint.__name__: route.endpoint + for route in router.routes[before:] + }["manual_cleanup"] + + +def _reference_database(monkeypatch, *, upload_id: str, gallery_hash: str = None): + from routes import upload_routes + + SessionLocal, engine, tmpfile = make_temp_sqlite(Base.metadata) + db = SessionLocal() + try: + db.add(DbSession( + id="session-1", + name="Cleanup regression", + endpoint_url="http://localhost", + model="test-model", + owner="alice", + )) + db.add(DbChatMessage( + id="message-1", + session_id="session-1", + role="user", + content=f"[Attachment: retained.png | id={upload_id} | mime=image/png]", + meta_data=json.dumps({ + "attachments": [{ + "id": upload_id, + "name": "retained.png", + "mime": "image/png", + "size": 8, + }] + }), + )) + if gallery_hash: + db.add(GalleryImage( + id="gallery-cleanup-reference", + filename="abcdef123456.png", + prompt="Chat upload", + owner="alice", + file_hash=gallery_hash, + )) + db.commit() + finally: + db.close() + + monkeypatch.setattr(upload_routes, "SessionLocal", SessionLocal) + return engine, tmpfile + + +def test_admin_cleanup_preserves_referenced_upload_and_reconciles_deleted_row( + tmp_path, + monkeypatch, +): + handler = _make_handler(tmp_path) + referenced_id = "a" * 32 + ".png" + unreferenced_id = "b" * 32 + ".txt" + gallery_id = "7" * 32 + ".png" + gallery_hash = "7" * 64 + paths = _seed_old_uploads(handler, [ + { + "id": referenced_id, + "hash": "1" * 64, + "mime": "image/png", + }, + { + "id": unreferenced_id, + "hash": "2" * 64, + "mime": "text/plain", + }, + { + "id": gallery_id, + "hash": gallery_hash, + "mime": "image/png", + }, + ]) + engine, tmpfile = _reference_database( + monkeypatch, + upload_id=referenced_id, + gallery_hash=gallery_hash, + ) + + try: + response = asyncio.run( + _manual_cleanup_endpoint(handler, monkeypatch)(_AdminRequest()) + ) + finally: + engine.dispose() + tmpfile.close() + try: + os.unlink(tmpfile.name) + except OSError: + pass + + assert response == {"status": "success", "files_cleaned": 1} + assert paths[referenced_id].is_file() + referenced_info = handler.get_upload_info(referenced_id) + assert referenced_info is not None + assert handler.resolve_upload(referenced_id, owner="alice") is not None + assert paths[gallery_id].is_file() + assert handler.get_upload_info(gallery_id) is not None + + assert not paths[unreferenced_id].exists() + assert handler.get_upload_info(unreferenced_id) is None + assert handler.resolve_upload(unreferenced_id, owner="alice") is None + + live_index = json.loads( + Path(handler.upload_dir, "uploads.json").read_text(encoding="utf-8") + ) + assert {info["id"] for info in live_index.values()} == { + referenced_id, + gallery_id, + } + backup_index = json.loads( + Path(handler.upload_dir, "uploads.json.bak").read_text(encoding="utf-8") + ) + assert {info["id"] for info in backup_index.values()} == { + referenced_id, + gallery_id, + } + + # Recovery must not resurrect the deliberately deleted row. + Path(handler.upload_dir, "uploads.json").write_text("{broken", encoding="utf-8") + handler._index_cache = None + assert handler.get_upload_info(unreferenced_id) is None + assert paths[referenced_id].parent.is_dir() + + +def test_cleanup_retains_upload_and_all_rows_when_index_rows_disagree(tmp_path): + handler = _make_handler(tmp_path) + upload_id = "c" * 32 + ".txt" + path = _seed_old_uploads(handler, [{ + "id": upload_id, + "hash": "1" * 64, + "mime": "text/plain", + "owner": "alice", + }])[upload_id] + index_path = Path(handler.upload_dir, "uploads.json") + index = json.loads(index_path.read_text(encoding="utf-8")) + alice_row = next(iter(index.values())) + index["bob:" + "2" * 64] = { + **alice_row, + "owner": "bob", + "hash": "2" * 64, + "checksum_sha256": "2" * 64, + } + index_path.write_text(json.dumps(index), encoding="utf-8") + handler._index_cache = None + + assert handler.cleanup_old_uploads(set(), set()) == 0 + assert path.is_file() + assert json.loads(index_path.read_text(encoding="utf-8")) == index + + +def test_cleanup_retains_lone_row_without_authoritative_lifecycle_metadata(tmp_path): + handler = _make_handler(tmp_path) + upload_id = "6" * 32 + ".txt" + path = _seed_old_uploads(handler, [{ + "id": upload_id, + "hash": "6" * 64, + "mime": "text/plain", + }])[upload_id] + index_path = Path(handler.upload_dir, "uploads.json") + index = json.loads(index_path.read_text(encoding="utf-8")) + row = next(iter(index.values())) + for field in ( + "owner", + "hash", + "checksum_sha256", + "uploaded_at", + "created_at", + "last_accessed", + ): + row.pop(field) + index_path.write_text(json.dumps(index), encoding="utf-8") + handler._index_cache = None + + assert handler.cleanup_old_uploads(set(), set()) == 0 + assert path.is_file() + assert json.loads(index_path.read_text(encoding="utf-8")) == index + + +def test_reservation_and_cleanup_are_serialized_without_dangling_references( + tmp_path, + monkeypatch, +): + # Writer wins: reservation holds the shared index lock, refreshes access, + # then cleanup observes the refreshed row and preserves the file. + writer_root = tmp_path / "writer-wins" + writer_root.mkdir() + writer_handler = _make_handler(writer_root) + upload_id = "2" * 32 + ".txt" + writer_path = _seed_old_uploads(writer_handler, [{ + "id": upload_id, + "hash": "2" * 64, + "mime": "text/plain", + }])[upload_id] + write_entered = threading.Event() + release_write = threading.Event() + real_atomic_write = writer_handler._atomic_write_json + + def blocking_reservation_write(path, data, *, sync_backup=False): + refreshed = any( + isinstance(row, dict) and row.get("last_accessed") != OLD_TIMESTAMP + for row in data.values() + ) + if sync_backup and refreshed and not write_entered.is_set(): + write_entered.set() + assert release_write.wait(5) + return real_atomic_write(path, data, sync_backup=sync_backup) + + monkeypatch.setattr(writer_handler, "_atomic_write_json", blocking_reservation_write) + with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool: + reserve_future = pool.submit( + writer_handler.reserve_upload, + upload_id, + owner="alice", + ) + assert write_entered.wait(5) + cleanup_future = pool.submit(writer_handler.cleanup_old_uploads, set(), set()) + release_write.set() + assert reserve_future.result(timeout=5) is not None + assert cleanup_future.result(timeout=5) == 0 + assert writer_path.is_file() + + # Cleanup wins: reservation cannot pass the same lock until the row and + # bytes are gone, then fails so a caller cannot commit a dangling reference. + cleanup_root = tmp_path / "cleanup-wins" + cleanup_root.mkdir() + cleanup_handler = _make_handler(cleanup_root) + cleanup_path = _seed_old_uploads(cleanup_handler, [{ + "id": upload_id, + "hash": "3" * 64, + "mime": "text/plain", + }])[upload_id] + remove_entered = threading.Event() + release_remove = threading.Event() + real_remove = os.remove + + def blocking_remove(candidate): + if os.path.realpath(candidate) == os.path.realpath(cleanup_path): + remove_entered.set() + assert release_remove.wait(5) + return real_remove(candidate) + + monkeypatch.setattr("src.upload_handler.os.remove", blocking_remove) + with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool: + cleanup_future = pool.submit(cleanup_handler.cleanup_old_uploads, set(), set()) + assert remove_entered.wait(5) + reserve_future = pool.submit( + cleanup_handler.reserve_upload, + upload_id, + owner="alice", + ) + release_remove.set() + assert cleanup_future.result(timeout=5) == 1 + assert reserve_future.result(timeout=5) is None + assert not cleanup_path.exists() + + +def test_admin_cleanup_reference_discovery_failure_returns_503_without_deleting( + tmp_path, + monkeypatch, +): + from routes import upload_routes + + handler = _make_handler(tmp_path) + upload_id = "d" * 32 + ".png" + path = _seed_old_uploads(handler, [ + {"id": upload_id, "hash": "4" * 64, "mime": "image/png"}, + ])[upload_id] + + def fail_reference_scan(): + raise RuntimeError("database unavailable") + + monkeypatch.setattr( + upload_routes, + "_collect_persisted_upload_references", + fail_reference_scan, + ) + endpoint = _manual_cleanup_endpoint(handler, monkeypatch) + + with pytest.raises(HTTPException) as exc: + asyncio.run(endpoint(_AdminRequest())) + + assert exc.value.status_code == 503 + assert path.is_file() + assert handler.get_upload_info(upload_id) is not None + + +def test_cleanup_restores_index_when_file_removal_fails(tmp_path, monkeypatch): + handler = _make_handler(tmp_path) + upload_id = "e" * 32 + ".txt" + path = _seed_old_uploads(handler, [ + { + "id": upload_id, + "hash": "5" * 64, + "mime": "text/plain", + }, + ])[upload_id] + + real_remove = os.remove + + def fail_target_remove(candidate): + if os.path.realpath(candidate) == os.path.realpath(path): + raise PermissionError("file is in use") + return real_remove(candidate) + + monkeypatch.setattr("src.upload_handler.os.remove", fail_target_remove) + + assert handler.cleanup_old_uploads(set(), set()) == 0 + assert path.is_file() + assert handler.get_upload_info(upload_id) is not None + assert any( + info["id"] == upload_id + for info in json.loads( + Path(handler.upload_dir, "uploads.json").read_text(encoding="utf-8") + ).values() + ) + assert any( + info["id"] == upload_id + for info in json.loads( + Path(handler.upload_dir, "uploads.json.bak").read_text(encoding="utf-8") + ).values() + ) + + +def test_admin_cleanup_with_corrupt_index_returns_503_and_fails_closed( + tmp_path, + monkeypatch, +): + from routes import upload_routes + + handler = _make_handler(tmp_path) + upload_id = "9" * 32 + ".png" + path = _seed_old_uploads(handler, [ + {"id": upload_id, "hash": "9" * 64, "mime": "image/png"}, + ])[upload_id] + Path(handler.upload_dir, "uploads.json").write_text( + '{"alice:broken": {', + encoding="utf-8", + ) + handler._index_cache = None + + monkeypatch.setattr( + upload_routes, + "_collect_persisted_upload_references", + lambda: (set(), set()), + ) + endpoint = _manual_cleanup_endpoint(handler, monkeypatch) + + with pytest.raises(HTTPException) as exc: + asyncio.run(endpoint(_AdminRequest())) + + assert exc.value.status_code == 503 + assert path.is_file() + + +def test_cleanup_with_missing_live_index_fails_closed(tmp_path): + handler = _make_handler(tmp_path) + upload_id = "8" * 32 + ".png" + dated_dir = Path(handler.upload_dir) / "2000" / "01" / "01" + dated_dir.mkdir(parents=True) + path = dated_dir / upload_id + path.write_bytes(b"unindexed bytes") + + with pytest.raises(UploadCleanupSafetyError): + handler.cleanup_old_uploads(set(), set()) + + assert path.is_file() + + +def test_reference_discovery_covers_all_durable_upload_stores( + monkeypatch, +): + from routes import upload_routes + + document_id = "f" * 32 + ".pdf" + version_id = "1" * 32 + ".pdf" + note_upload_id = "3" * 32 + ".png" + note_color_id = "2" * 32 + ".png" + calendar_upload_id = "4" * 32 + ".png" + event_upload_id = "5" * 32 + ".png" + event_description_id = "7" * 32 + ".txt" + event_location_id = "8" * 32 + ".png" + gallery_hash = "6" * 64 + SessionLocal, engine, tmpfile = make_temp_sqlite(Base.metadata) + db = SessionLocal() + try: + db.add(DbSession( + id="session-2", + name="Reference sources", + endpoint_url="http://localhost", + model="test-model", + owner="alice", + )) + db.add(Document( + id="document-1", + session_id="session-2", + title="PDF", + current_content=f'', + owner="alice", + )) + db.add(DocumentVersion( + id="version-1", + document_id="document-1", + version_number=1, + content=f'', + )) + db.add(GalleryImage( + id="gallery-1", + # Gallery filenames are normally generated 12-hex names, so this + # record proves retention comes from its stored content hash. + filename="abcdef123456.png", + prompt="Chat upload", + owner="alice", + file_hash=gallery_hash, + )) + db.add(Note( + id="note-1", + owner="alice", + title="Photo note", + image_url=f"/api/upload/{note_upload_id}", + color=f"odysseus://attachment/{note_color_id}", + )) + db.add(CalendarCal( + id="calendar-1", + owner="alice", + name="Personal", + color=f"/api/upload/{calendar_upload_id}", + )) + db.add(CalendarEvent( + uid="event-1", + calendar_id="calendar-1", + summary="Photo event", + dtstart=datetime(2026, 7, 10, 12, 0), + dtend=datetime(2026, 7, 10, 13, 0), + color=f"/api/upload/{event_upload_id}", + description=f"Notes: odysseus://attachment/{event_description_id}", + location=f"/api/upload/{event_location_id}", + )) + db.commit() + finally: + db.close() + + monkeypatch.setattr(upload_routes, "SessionLocal", SessionLocal) + try: + referenced_ids, referenced_hashes = ( + upload_routes._collect_persisted_upload_references() + ) + finally: + engine.dispose() + tmpfile.close() + try: + os.unlink(tmpfile.name) + except OSError: + pass + + assert { + document_id, + version_id, + note_upload_id, + note_color_id, + calendar_upload_id, + event_upload_id, + event_description_id, + event_location_id, + } <= referenced_ids + assert gallery_hash in referenced_hashes + + +def test_write_reservation_extracts_only_explicit_internal_references(): + upload_id = "a" * 32 + ".png" + checksum_like_text = "b" * 32 + + assert extract_internal_upload_ids(checksum_like_text) == set() + assert extract_internal_upload_ids(f"sha={checksum_like_text}") == set() + assert extract_internal_upload_ids({ + "image": f"/api/upload/{upload_id}", + "nested": [f"odysseus://attachment/{upload_id}"], + }) == {upload_id} + assert extract_internal_upload_ids( + f'' + ) == {upload_id} + assert extract_internal_upload_ids( + f"[Attachment: photo.png | id={upload_id} | mime=image/png]" + ) == {upload_id} + extensionless_id = "c" * 32 + assert extract_internal_upload_ids( + f"See /api/upload/{extensionless_id}. Then continue." + ) == {extensionless_id} + assert extract_internal_upload_ids( + f"Attachment: odysseus://attachment/{extensionless_id}: ready" + ) == {extensionless_id} + assert extract_internal_upload_ids(f"/api/upload/{upload_id}/extra") == set() + + +def test_reservation_never_uses_admin_override(tmp_path): + handler = _make_handler(tmp_path) + upload_id = "c" * 32 + ".txt" + _seed_old_uploads(handler, [{ + "id": upload_id, + "hash": "c" * 64, + "mime": "text/plain", + "owner": "alice", + }]) + + assert reserve_upload_references( + handler, + "alice", + f"odysseus://attachment/{upload_id}", + ) is None + assert reserve_upload_references( + handler, + "admin", + f"odysseus://attachment/{upload_id}", + ) == upload_id + assert handler.reserve_upload( + upload_id, + owner="admin", + auth_manager=_AdminAuth(), + allow_admin=False, + ) is None + assert reserve_message_upload_references( + handler, + "admin", + "legacy attachment metadata", + {"attachments": [{"id": upload_id, "name": "owned.txt"}]}, + ) == upload_id + + +def test_remaining_durable_writers_reserve_before_commit(monkeypatch): + import core.database as database + import core.session_manager as session_manager_module + import src.database as legacy_database + from core.models import ChatMessage + from core.session_manager import SessionManager + from src import tool_utils + from src.agent_tools.document_tools import EditDocumentTool + from src.tools.calendar import do_manage_calendar + from src.tools.notes import do_manage_notes + + upload_id = "6" * 32 + ".png" + + class RejectingHandler: + def reserve_upload(self, _candidate, **_kwargs): + return None + + handler = RejectingHandler() + monkeypatch.setattr(tool_utils, "_upload_handler", handler) + + SessionLocal, engine, tmpfile = make_temp_sqlite(Base.metadata) + monkeypatch.setattr(database, "SessionLocal", SessionLocal) + monkeypatch.setattr(legacy_database, "SessionLocal", SessionLocal) + monkeypatch.setattr(legacy_database, "Document", Document, raising=False) + monkeypatch.setattr( + legacy_database, + "DocumentVersion", + DocumentVersion, + raising=False, + ) + monkeypatch.setattr(session_manager_module, "SessionLocal", SessionLocal) + db = SessionLocal() + try: + db.add(DbSession( + id="writer-session", + name="Writer coverage", + endpoint_url="http://localhost", + model="test-model", + owner="alice", + )) + db.add(Document( + id="email-document", + session_id="writer-session", + title="New Email", + language="email", + current_content="To: team@example.test\nSubject: Status\n---\nOld body", + version_count=1, + owner="alice", + )) + db.commit() + + manager = SessionManager() + manager.upload_handler = handler + manager._persist_message( + "writer-session", + ChatMessage( + "user", + "attachment", + metadata={"attachments": [{"id": upload_id}]}, + ), + ) + + document_result = asyncio.run(EditDocumentTool().execute( + "<<>>\n\n<<>>\n" + f"See /api/upload/{upload_id}\n<<>>", + {"doc_id": "email-document", "owner": "alice"}, + )) + assert document_result["exit_code"] == 1 + assert "no longer available" in document_result["error"] + + calendar_result = asyncio.run(do_manage_calendar( + json.dumps({ + "action": "create_event", + "summary": "Attachment review", + "dtstart": "2026-07-12T12:00:00", + "description": f"See /api/upload/{upload_id}", + }), + owner="alice", + )) + assert calendar_result["exit_code"] == 1 + assert "no longer available" in calendar_result["error"] + + note_result = asyncio.run(do_manage_notes( + json.dumps({ + "action": "add", + "title": "Attachment note", + "content": f"See /api/upload/{upload_id}", + }), + owner="alice", + )) + assert note_result["exit_code"] == 1 + assert "no longer available" in note_result["error"] + + verify = SessionLocal() + try: + assert verify.query(DbChatMessage).count() == 0 + stored_doc = verify.query(Document).filter(Document.id == "email-document").one() + assert stored_doc.current_content.endswith("Old body") + assert verify.query(CalendarEvent).count() == 0 + assert verify.query(Note).count() == 0 + finally: + verify.close() + finally: + db.close() + engine.dispose() + tmpfile.close() + try: + os.unlink(tmpfile.name) + except OSError: + pass + + +def test_note_calendar_and_document_routes_reserve_before_database_writes(monkeypatch): + from routes.calendar_routes import EventCreate, setup_calendar_routes + from routes import document_routes + from routes.document_helpers import DocumentCreate + from routes.note_routes import NoteCreate, setup_note_routes + from src import auth_helpers + + upload_id = "d" * 32 + ".png" + + class RejectingHandler: + def __init__(self): + self.calls = [] + + def reserve_upload(self, candidate, **kwargs): + self.calls.append((candidate, kwargs)) + return None + + request = SimpleNamespace( + state=SimpleNamespace(current_user="alice", api_token=False), + app=SimpleNamespace(state=SimpleNamespace()), + ) + + note_handler = RejectingHandler() + note_router = setup_note_routes(upload_handler=note_handler) + create_note = next( + route.endpoint for route in note_router.routes + if route.endpoint.__name__ == "create_note" + ) + with pytest.raises(HTTPException) as note_error: + create_note( + request, + NoteCreate(image_url=f"/api/upload/{upload_id}"), + ) + assert note_error.value.status_code == 409 + assert note_handler.calls == [ + (upload_id, {"owner": "alice", "allow_admin": False}) + ] + + calendar_handler = RejectingHandler() + calendar_router = setup_calendar_routes(upload_handler=calendar_handler) + create_event = next( + route.endpoint for route in calendar_router.routes + if route.endpoint.__name__ == "create_event" + ) + with pytest.raises(HTTPException) as calendar_error: + asyncio.run(create_event( + request, + EventCreate( + summary="Photo", + dtstart="2026-07-10T12:00:00", + color=f"odysseus://attachment/{upload_id}", + ), + )) + assert calendar_error.value.status_code == 409 + assert calendar_handler.calls == [ + (upload_id, {"owner": "alice", "allow_admin": False}) + ] + + class EmptyDb: + @staticmethod + def close(): + return None + + document_handler = RejectingHandler() + monkeypatch.setattr(document_routes, "SessionLocal", EmptyDb) + monkeypatch.setattr( + auth_helpers, + "require_privilege", + lambda _request, _privilege: "alice", + ) + document_router = document_routes.setup_document_routes( + SimpleNamespace(), + document_handler, + ) + create_document = next( + route.endpoint for route in document_router.routes + if route.endpoint.__name__ == "create_document" + ) + with pytest.raises(HTTPException) as document_error: + asyncio.run(create_document( + request, + DocumentCreate( + language="markdown", + content=f"![image](/api/upload/{upload_id})", + ), + )) + assert document_error.value.status_code == 409 + assert document_handler.calls == [ + (upload_id, {"owner": "alice", "allow_admin": False}) + ]