Squash Odysseus development history

core/session_manager.py

@@ -14,12 +14,54 @@ import logging
 from datetime import datetime, timezone, timedelta
 from typing import Dict, Optional
 
-from .database import Session as DbSession, ChatMessage as DbChatMessage, Document as DbDocument, SessionLocal
+from sqlalchemy import func
+
+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
 
 logger = logging.getLogger(__name__)
 
 
+def _message_timestamp_iso(value: Optional[datetime]) -> Optional[str]:
+    """Return a stable ISO timestamp for chat message metadata."""
+    if not value:
+        return None
+    if value.tzinfo is None:
+        value = value.replace(tzinfo=timezone.utc)
+    return value.isoformat().replace("+00:00", "Z")
+
+
+def _parse_msg_content(raw):
+    """Parse message content from DB — deserialises JSON arrays back to lists
+    (multimodal content with image/audio attachments)."""
+    if isinstance(raw, list):
+        return raw
+    if isinstance(raw, str) and raw.startswith('[{') and '"type"' in raw:
+        try:
+            parsed = json.loads(raw)
+            # Only treat as serialized multimodal content when EVERY element is
+            # a dict whose "type" is a recognized content-block kind. Otherwise a
+            # plain text message that merely *looks* like a JSON array of objects
+            # (e.g. a user pasting an API schema/sample with a "type" field) was
+            # silently parsed back into a list, destroying the original string.
+            _BLOCK_TYPES = {
+                "text", "image", "image_url", "audio", "input_audio",
+                "input_image", "document", "file",
+            }
+            if (isinstance(parsed, list) and parsed
+                    and all(isinstance(p, dict) and p.get("type") in _BLOCK_TYPES
+                            for p in parsed)):
+                return parsed
+        except (json.JSONDecodeError, ValueError):
+            pass
+    return raw
+
+
 class SessionManager:
     """
     Manages chat sessions with database persistence.
@@ -34,6 +76,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()
 
     # ------------------------------------------------------------------
@@ -51,14 +94,28 @@ class SessionManager:
         try:
             db_sessions = db.query(DbSession).filter(
                 DbSession.archived == False,
-                DbSession.message_count > 0,
+                DbSession.messages.any(),
             ).order_by(DbSession.last_accessed.desc()).limit(100).all()
 
+            # message_count is derived metadata and can drift after interrupted
+            # or legacy writes. Count only the bounded discovery set so startup
+            # remains metadata-only while lazy hydration sees an authoritative
+            # positive count for every discovered non-empty session.
+            message_counts = {}
+            if db_sessions:
+                message_counts = dict(
+                    db.query(DbChatMessage.session_id, func.count(DbChatMessage.id))
+                    .filter(DbChatMessage.session_id.in_([row.id for row in db_sessions]))
+                    .group_by(DbChatMessage.session_id)
+                    .all()
+                )
+
             loaded_count = 0
             for db_session in db_sessions:
                 try:
                     session = self._db_to_session_meta(db_session)
                     if session is not None:
+                        session.message_count = message_counts[db_session.id]
                         self.sessions[db_session.id] = session
                         loaded_count += 1
                 except Exception as e:
@@ -93,6 +150,12 @@ class SessionManager:
             history=[],
             owner=getattr(db_session, "owner", None),
             is_important=getattr(db_session, "is_important", False) or False,
+            memory_extraction_enabled=getattr(db_session, "memory_extraction_enabled", True) is not False,
+            skill_injection_enabled=getattr(db_session, "skill_injection_enabled", True) is not False,
+            thinking_mode=getattr(db_session, "thinking_mode", "") or "off",
+            temperature_override=getattr(db_session, "temperature_override", None),
+            max_tokens_override=getattr(db_session, "max_tokens_override", None),
+            cwd=getattr(db_session, "cwd", None) or None,
         )
         session.message_count = getattr(db_session, "message_count", 0) or 0
         return session
@@ -107,9 +170,10 @@ class SessionManager:
                 meta = json.loads(db_msg.meta_data) if db_msg.meta_data else {}
                 if meta is None: meta = {}
                 meta['_db_id'] = db_msg.id
+                meta.setdefault('timestamp', _message_timestamp_iso(db_msg.timestamp))
                 history.append(ChatMessage(
                     role=db_msg.role,
-                    content=db_msg.content,
+                    content=_parse_msg_content(db_msg.content),
                     metadata=meta,
                 ))
         else:
@@ -121,9 +185,10 @@ class SessionManager:
                 meta = json.loads(db_msg.meta_data) if db_msg.meta_data else {}
                 if meta is None: meta = {}
                 meta['_db_id'] = db_msg.id
+                meta.setdefault('timestamp', _message_timestamp_iso(db_msg.timestamp))
                 history.append(ChatMessage(
                     role=db_msg.role,
-                    content=db_msg.content,
+                    content=_parse_msg_content(db_msg.content),
                     metadata=meta,
                 ))
 
@@ -149,9 +214,20 @@ class SessionManager:
             history=history,
             owner=getattr(db_session, 'owner', None),
             is_important=getattr(db_session, 'is_important', False) or False,
+            memory_extraction_enabled=getattr(db_session, 'memory_extraction_enabled', True) is not False,
+            skill_injection_enabled=getattr(db_session, 'skill_injection_enabled', True) is not False,
+            thinking_mode=getattr(db_session, "thinking_mode", "") or "off",
+            temperature_override=getattr(db_session, "temperature_override", None),
+            max_tokens_override=getattr(db_session, "max_tokens_override", None),
+            cwd=getattr(db_session, "cwd", None) or None,
         )
 
-        session.message_count = getattr(db_session, 'message_count', len(history))
+        # The rows just loaded are the whole transcript, so they — not the
+        # denormalized sessions.message_count column — are the truth for this
+        # cached object. get_session's hydration gate compares against this
+        # number; seeding it from a drifted column would ask for a reload that
+        # can never close the gap.
+        session.message_count = len(history)
         return session
 
     # ------------------------------------------------------------------
@@ -162,12 +238,17 @@ class SessionManager:
         """
         Add a message to a session and persist to database.
 
+        Updates the authoritative history list and persists through this
+        manager directly so tests and temporary managers do not depend on the
+        process-wide session-manager singleton.
+
         Args:
             session_id: Session ID
             message: ChatMessage to add
         """
         session = self.get_session(session_id)
         session.history.append(message)
+        session._history = session.history
         session.message_count = len(session.history)
 
         self._persist_message(session_id, message)
@@ -176,31 +257,59 @@ class SessionManager:
         """Persist a single message to the database."""
         db = SessionLocal()
         try:
+            db_session = db.query(DbSession).filter(DbSession.id == session_id).first()
+            if db_session is None:
+                # A stream/tool callback can outlive a session delete. Do not
+                # create a chat_messages row with no parent session; also drop
+                # any stale cached session so later writes fail closed too.
+                self.sessions.pop(session_id, None)
+                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 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,
                 role=message.role,
-                content=message.content,
-                meta_data=json.dumps(message.metadata) if message.metadata else None
+                content=_content,
+                meta_data=json.dumps(message.metadata) if message.metadata else None,
+                timestamp=msg_time,
             )
             db.add(db_message)
 
-            db_session = db.query(DbSession).filter(DbSession.id == session_id).first()
-            if db_session:
-                db_session.message_count = len(self.sessions.get(session_id, {}).history) if session_id in self.sessions else 0
-                _now = datetime.now(timezone.utc)
-                db_session.last_accessed = _now
-                # Clean "last conversation" timestamp — only bumped here on a
-                # real message persist, so it powers an accurate "Last active"
-                # sort that ignores renames / model swaps / mere opens.
-                db_session.last_message_at = _now
+            if session_id in self.sessions:
+                db_session.message_count = len(self.sessions[session_id].history)
+            else:
+                db_session.message_count = 0
+            _now = datetime.now(timezone.utc)
+            db_session.last_accessed = _now
+            # Clean "last conversation" timestamp — only bumped here on a
+            # real message persist, so it powers an accurate "Last active"
+            # sort that ignores renames / model swaps / mere opens.
+            db_session.last_message_at = _now
 
             db.commit()
 
             # Store DB ID on the in-memory message for edit/delete by ID
-            if message.metadata is None:
-                message.metadata = {}
             message.metadata['_db_id'] = msg_id
 
             logger.debug(f"Persisted message to session {session_id}")
@@ -231,13 +340,17 @@ class SessionManager:
 
             db_session = db.query(DbSession).filter(DbSession.id == session_id).first()
             if db_session:
-                db_session.message_count = keep_count
+                # keep_count can exceed the real message total (e.g. the AI tool
+                # defaults to keep_count=10 on a short session); message_count must
+                # track the rows that actually remain, not the requested cap.
+                db_session.message_count = min(keep_count, len(db_messages))
                 db_session.updated_at = datetime.now(timezone.utc)
 
             db.commit()
 
             # Update in-memory
             session.history = session.history[:keep_count]
+            session._history = session.history
 
             logger.info(f"Truncated session {session_id} to {keep_count} messages")
             return True
@@ -254,6 +367,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):
@@ -262,7 +397,9 @@ class SessionManager:
                     id=msg_id,
                     session_id=session_id,
                     role=message.role,
-                    content=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),
                 )
@@ -271,15 +408,14 @@ 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)
+            session._history = session.history
             session.message_count = len(messages)
             logger.info("Replaced session %s history with %d messages", session_id, len(messages))
             return True
@@ -295,24 +431,85 @@ class SessionManager:
     # ------------------------------------------------------------------
 
     def get_session(self, session_id: str) -> Session:
-        """Get a session by ID, loading from DB if needed.
+        """Get a session by ID, loading complete DB history when needed.
 
-        Sessions seeded by `load_sessions` start with empty history. The
-        first read here hydrates them with the message rows.
+        Sessions seeded by ``load_sessions`` start with empty history, and a
+        cached session can also become partially stale. Refresh metadata first,
+        then hydrate whenever the cached transcript is short of the stored rows.
+        Model-send routes enter through this method before building context,
+        while paginated display history reads SQLite directly.
+
+        The gate compares against ``sync_session_metadata``'s reconciled count
+        (the real ``chat_messages`` total), never the denormalized column, so a
+        hydrate always closes the gap and the next read is a cache hit.
         """
         if session_id not in self.sessions:
             self._load_session_from_db(session_id)
-        else:
-            cached = self.sessions[session_id]
-            # Lazy hydrate: metadata-only entries get their messages on first read.
-            if not cached.history and getattr(cached, "message_count", 0) > 0:
-                self._load_session_from_db(session_id)
+
+        # Keep model/endpoint metadata fresh. Endpoint deletion can clear the
+        # DB row while a session object is still cached in RAM. Refreshing first
+        # also exposes the authoritative message count before completeness is
+        # checked.
+        self.sync_session_metadata(session_id)
+
+        cached = self.sessions[session_id]
+        cached_count = len(cached.history or [])
+        stored_count = int(getattr(cached, "message_count", 0) or 0)
+        if cached_count < stored_count:
+            self._load_session_from_db(session_id)
 
         # Update last_accessed
         self._touch_session(session_id)
 
         return self.sessions[session_id]
 
+    def sync_session_metadata(self, session_id: str) -> bool:
+        """Refresh non-message session fields from the DB into the cached object.
+
+        ``message_count`` is reconciled against the real ``chat_messages`` rows
+        rather than copied from the denormalized ``sessions.message_count``
+        column. That column drifts in normal operation — ``_persist_message``
+        swallows a failed insert but ``add_message`` has already appended in
+        memory, so the next successful persist writes rows+1, and a persist for
+        an uncached session writes 0. Hydration keys off this number: a
+        drifted-high column would reload the whole transcript on every warm
+        read, and a drifted-low one would leave the model a truncated one.
+        """
+        session = self.sessions.get(session_id)
+        if session is None:
+            return False
+        db = SessionLocal()
+        try:
+            db_session = db.query(DbSession).filter(DbSession.id == session_id).first()
+            if db_session is None:
+                return False
+            headers = db_session.headers
+            if isinstance(headers, str):
+                try:
+                    headers = json.loads(headers)
+                except json.JSONDecodeError:
+                    headers = {}
+            session.name = db_session.name
+            session.endpoint_url = db_session.endpoint_url or ""
+            session.model = db_session.model or ""
+            session.headers = headers or {}
+            session.rag = db_session.rag
+            session.archived = db_session.archived
+            session.owner = getattr(db_session, "owner", None)
+            session.is_important = getattr(db_session, "is_important", False) or False
+            session.cwd = getattr(db_session, "cwd", None) or None
+            session.message_count = (
+                db.query(DbChatMessage)
+                .filter(DbChatMessage.session_id == session_id)
+                .count()
+            )
+            return True
+        except Exception as e:
+            logger.error(f"Error syncing session metadata {session_id}: {e}")
+            return False
+        finally:
+            db.close()
+
     def _load_session_from_db(self, session_id: str):
         """Hydrate a single session (with messages) from the database."""
         db = SessionLocal()
@@ -361,9 +558,12 @@ class SessionManager:
         endpoint_url: str,
         model: str,
         rag: bool = False,
-        owner: str = None
+        owner: str = None,
+        cwd: str = None,
+        headers: Optional[Dict[str, str]] = None,
     ) -> Session:
         """Create a new session and save to database."""
+        session_headers = dict(headers or {})
         db = SessionLocal()
         try:
             db_session = DbSession(
@@ -372,8 +572,9 @@ class SessionManager:
                 endpoint_url=endpoint_url,
                 model=model,
                 rag=rag,
-                headers={},
+                headers=session_headers,
                 owner=owner,
+                cwd=cwd or None,
                 created_at=datetime.now(timezone.utc),
                 updated_at=datetime.now(timezone.utc)
             )
@@ -386,8 +587,9 @@ class SessionManager:
                 endpoint_url=endpoint_url,
                 model=model,
                 rag=rag,
-                headers={},
+                headers=session_headers,
                 owner=owner,
+                cwd=cwd or None,
             )
 
             self.sessions[session_id] = session
@@ -404,6 +606,12 @@ class SessionManager:
         """Permanently delete a session and all its messages."""
         db = SessionLocal()
         try:
+            try:
+                from src.session_image_cleanup import cleanup_session_images
+                cleanup_session_images(session_id, db=db)
+            except Exception as e:
+                logger.warning(f"Image cleanup failed while deleting session {session_id}: {e}")
+
             # Detach documents so they survive as orphans in the library
             db.query(DbDocument).filter(DbDocument.session_id == session_id).update(
                 {DbDocument.session_id: None}, synchronize_session=False
@@ -416,11 +624,17 @@ class SessionManager:
             db_session = db.query(DbSession).filter(DbSession.id == session_id).first()
             if db_session:
                 db.delete(db_session)
+
+            # Drop the in-memory copy even when there is no DB row. A "ghost"
+            # session lives only here (never persisted, or its row was removed
+            # out-of-band); without this it can never be cleared and keeps
+            # 404ing on every operation (issue #1044).
+            removed_in_memory = self.sessions.pop(session_id, None) is not None
+
+            if db_session or removed_in_memory:
+                # Commit the document-detach / message-delete above (a no-op when
+                # the ghost had no rows) together with the session delete.
                 db.commit()
-
-                if session_id in self.sessions:
-                    del self.sessions[session_id]
-
                 logger.info(f"Deleted session {session_id}")
                 return True
             return False
@@ -513,24 +727,52 @@ class SessionManager:
     def save_sessions(self):
         """No-op for DB compatibility."""
 
+    def ensure_task_session(self, session_id: str, name: str, endpoint_url: str, model: str, owner: str = None, task: object = None) -> Session:
+        """Create a task session if it doesn't exist, or return the existing one.
+
+        Unlike create_session, this checks the cache first and does NOT
+        overwrite an existing in-memory session. The task scheduler must
+        use this instead of direct dict assignment.
+        """
+        if session_id in self.sessions:
+            return self.sessions[session_id]
+
+        session = self.create_session(session_id, name, endpoint_url, model, owner=owner)
+        if task is not None:
+            task.session_id = session_id
+        return session
+
     # ------------------------------------------------------------------
     # Cleanup
     # ------------------------------------------------------------------
 
-    def cleanup_empty_sessions(self, auto_archive_days: int = 30) -> dict:
-        """Clean up empty and old sessions."""
+    def cleanup_empty_sessions(self, auto_archive_days: int = 30, min_age_hours: int = 1) -> dict:
+        """Clean up empty and old sessions.
+
+        Args:
+            auto_archive_days: Age in days before non-important sessions are archived.
+            min_age_hours: Minimum age in hours before an empty session can be deleted.
+                          Prevents deleting sessions that were just created.
+        """
         db = SessionLocal()
         stats = {'deleted_empty': 0, 'archived_old': 0, 'total_checked': 0}
 
         try:
             all_sessions = db.query(DbSession).all()
-            cutoff_date = datetime.now(timezone.utc) - timedelta(days=auto_archive_days)
+            cutoff_date = utcnow_naive() - timedelta(days=auto_archive_days)
+            min_age = utcnow_naive() - timedelta(hours=min_age_hours)
 
             for db_session in all_sessions:
                 stats['total_checked'] += 1
 
-                # Delete empty sessions
+                # Delete empty sessions only if older than min_age_hours
                 if db_session.message_count == 0:
+                    if db_session.created_at is not None:
+                        created = db_session.created_at
+                        if created.tzinfo is None:
+                            created = created.replace(tzinfo=timezone.utc)
+                        if created > min_age:
+                            continue  # Too young to delete
                     if db_session.id in self.sessions:
                         del self.sessions[db_session.id]
                     db.delete(db_session)