fix: session context drifting — messages leaking between chats (#135) (#267)

* docs: add implementation plan for fixing chat context drifting (#135)

* fix: make Session.history immutable + fix {}.history crash

- Session.history now exposes a COPY of the internal _history list
- add_message() replaces history with a fresh copy each time
- get_context_messages() derives from _history directly
- replace_messages() updates both _history and history
- truncate_messages() updates both _history and history
- _persist_message() line 207: fixed {}.history fallback crash
- Added 11 tests for session isolation and edge cases

Addresses #135 root cause #1: shared mutable references

* fix: task scheduler uses SessionManager methods instead of overwriting sessions

- Added ensure_task_session() to SessionManager (checks cache first)
- Task scheduler now uses ensure_task_session() instead of direct dict assignment
- Task scheduler now uses SessionManager.add_message() for message persistence
- Removed direct sess_obj.history.append() that was silently losing data

Addresses #135 root causes #2 and #3

* fix: add age guard to cleanup_empty_sessions — don't delete sessions <1h old

Prevents the cleanup task from deleting sessions that were just created
and haven't received any messages yet (message_count == 0).

Addresses #135 root cause #5

* test: comprehensive session isolation tests (10/10 passing)

* refactor: consolidate _session_manager into singleton pattern

- Added set_session_manager_instance / get_session_manager_instance to core/models
- kept backward-compat aliases (set_session_manager, get_session_manager)
- session_manager.py re-exports the singleton functions
- ai_interaction.set_session_manager now syncs with the core singleton
- context_compactor uses get_session_manager_instance() instead of getattr hack
- app.py initializes the singleton once

Addresses #135 root cause #4: fragile global wiring

* test: add concurrent session isolation integration tests

Verifies:
- Concurrent add_message to different sessions doesn't cross-contaminate
- Rapid parallel writes maintain isolation
- Read-write concurrent access is safe

All 3 async tests pass, proving the immutable history fix works under concurrency

* fix: pre-import core.models in conftest to prevent test pollution

test_agent_loop.py stubs sys.modules['core.models'] = MagicMock() at
module level during collection. Any test collected after it imports
Session as a MagicMock. Pre-importing core.models in conftest.py
before test_agent_loop.py's module-level code runs prevents this.

* fix: make .history authoritative mutable list, address PR review

Per review feedback: keep .history as the authoritative mutable list so
existing code doing .history.pop(), .history = [...], etc. still works.
Fix the cross-contamination bug by ensuring __post_init__() gives each
Session its OWN unique history list (never shared).

Changes:
- core/models.py: .history IS the authoritative list. _history aliases it.
  Each Session gets its own list in __post_init__.
- core/session_manager.py: add_message() delegates to Session.add_message()
  instead of appending directly — no double-append, single source of truth.
- tests/test_session_manager.py: updated test to reflect that .history
  references see new messages (same list, not a snapshot).
- docs/plans/2026-06-01-fix-chat-context-drifting.md: removed (not for
  shipping — useful design context but too much process/doc to ship).

All 272 tests pass (3 pre-existing failures unrelated).

* Fix session manager message persistence

* Fix session history alias regressions

* Fix session history aliasing and task delivery
This commit is contained in:
Joshua Valderrama
2026-06-09 21:12:52 +08:00
committed by GitHub
parent c3fcaf15b7
commit 35b4dd2824
12 changed files with 542 additions and 53 deletions
+45 -4
View File
@@ -17,6 +17,9 @@ 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
# Re-export singleton accessors from models for convenience
from .models import set_session_manager_instance, get_session_manager_instance
logger = logging.getLogger(__name__)
@@ -188,12 +191,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)
@@ -232,7 +240,10 @@ class SessionManager:
)
db.add(db_message)
db_session.message_count = len(self.sessions.get(session_id, {}).history) if session_id in self.sessions else 0
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
@@ -283,6 +294,7 @@ class SessionManager:
# 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
@@ -333,6 +345,7 @@ class SessionManager:
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
@@ -608,24 +621,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 = 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)