Merge PR #3440: fix: Automatically clean up unsupported or expired MSTeams session

fix: Automatically clean up unsupported or expired MSTeams session

nanobot/channels/msteams.py

@@ -15,12 +15,21 @@ import asyncio
 import html
 import importlib.util
 import json
+import os
 import re
+import tempfile
 import threading
 import time
+from contextlib import contextmanager
 from dataclasses import dataclass
 from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
 from typing import TYPE_CHECKING, Any
+from urllib.parse import urlparse
+
+try:  # pragma: no cover - Windows fallback path
+    import fcntl
+except ImportError:  # pragma: no cover
+    fcntl = None
 
 import httpx
 from loguru import logger
@@ -43,6 +52,13 @@ if TYPE_CHECKING:
 if MSTEAMS_AVAILABLE:
     import jwt
 
+MSTEAMS_REF_TTL_DAYS = 30
+MSTEAMS_REF_TTL_S = MSTEAMS_REF_TTL_DAYS * 24 * 60 * 60
+MSTEAMS_WEBCHAT_HOST = "webchat.botframework.com"
+MSTEAMS_REF_META_FILENAME = "msteams_conversations_meta.json"
+MSTEAMS_REF_LOCK_FILENAME = "msteams_conversations.lock"
+MSTEAMS_REF_TOUCH_INTERVAL_S = 300
+
 
 class MSTeamsConfig(Base):
     """Microsoft Teams channel configuration."""
@@ -58,6 +74,10 @@ class MSTeamsConfig(Base):
     reply_in_thread: bool = True
     mention_only_response: str = "Hi — what can I help with?"
     validate_inbound_auth: bool = True
+    ref_ttl_days: int = Field(default=MSTEAMS_REF_TTL_DAYS, ge=1)
+    prune_web_chat_refs: bool = True
+    prune_non_personal_refs: bool = True
+    ref_touch_interval_s: int = Field(default=MSTEAMS_REF_TOUCH_INTERVAL_S, ge=0)
 
 
 @dataclass
@@ -103,7 +123,13 @@ class MSTeamsChannel(BaseChannel):
         self._botframework_jwks_expires_at: float = 0.0
         self._refs_path = get_workspace_path() / "state" / "msteams_conversations.json"
         self._refs_path.parent.mkdir(parents=True, exist_ok=True)
+        self._refs_meta_path = self._refs_path.parent / MSTEAMS_REF_META_FILENAME
+        self._refs_lock_path = self._refs_path.parent / MSTEAMS_REF_LOCK_FILENAME
+        self._refs_guard = threading.RLock()
         self._conversation_refs: dict[str, ConversationRef] = self._load_refs()
+        with self._refs_guard:
+            if self._prune_conversation_refs():
+                self._save_refs_locked(prune=True)
 
     async def start(self) -> None:
         """Start the Teams webhook listener."""
@@ -236,6 +262,7 @@ class MSTeamsChannel(BaseChannel):
             resp = await self._http.post(base_url, headers=headers, json=payload)
             resp.raise_for_status()
             logger.info("MSTeams message sent to {}", ref.conversation_id)
+            self._touch_conversation_ref(str(msg.chat_id), persist=True)
         except Exception as e:
             logger.error("MSTeams send failed: {}", e)
             raise
@@ -282,17 +309,17 @@ class MSTeamsChannel(BaseChannel):
             )
             return
 
-        self._conversation_refs[conversation_id] = ConversationRef(
-            service_url=service_url,
-            conversation_id=conversation_id,
-            bot_id=str(recipient.get("id") or "") or None,
-            activity_id=activity_id or None,
-            conversation_type=conversation_type or None,
-            tenant_id=str((channel_data.get("tenant") or {}).get("id") or "") or None,
-            updated_at=time.time(),
-        )
-
-        self._save_refs()
+        with self._refs_guard:
+            self._conversation_refs[conversation_id] = ConversationRef(
+                service_url=service_url,
+                conversation_id=conversation_id,
+                bot_id=str(recipient.get("id") or "") or None,
+                activity_id=activity_id or None,
+                conversation_type=conversation_type or None,
+                tenant_id=str((channel_data.get("tenant") or {}).get("id") or "") or None,
+                updated_at=time.time(),
+            )
+            self._save_refs_locked()
 
         await self._handle_message(
             sender_id=sender_id,
@@ -487,61 +514,242 @@ class MSTeamsChannel(BaseChannel):
         self._botframework_jwks_expires_at = now + 3600
         return self._botframework_jwks
 
+    @staticmethod
+    def _safe_float(value: Any) -> float | None:
+        try:
+            out = float(value)
+            if out > 0:
+                return out
+        except (TypeError, ValueError):
+            return None
+        return None
+
+    def _normalize_ref_record(self, value: Any) -> ConversationRef | None:
+        """Normalize a stored ref record from legacy/current schema."""
+        if not isinstance(value, dict):
+            return None
+        service_url = str(value.get("service_url") or "").strip()
+        conversation_id = str(value.get("conversation_id") or "").strip()
+        if not service_url or not conversation_id:
+            return None
+        return ConversationRef(
+            service_url=service_url,
+            conversation_id=conversation_id,
+            bot_id=str(value.get("bot_id") or "") or None,
+            activity_id=str(value.get("activity_id") or "") or None,
+            conversation_type=str(value.get("conversation_type") or "") or None,
+            tenant_id=str(value.get("tenant_id") or "") or None,
+            updated_at=self._safe_float(value.get("updated_at")),
+        )
+
+    def _load_refs_raw(self) -> tuple[dict[str, Any], dict[str, Any], bool]:
+        """Load raw refs/main+meta JSON payloads."""
+        main_data: dict[str, Any] = {}
+        meta_data: dict[str, Any] = {}
+        meta_exists = self._refs_meta_path.exists()
+
+        if self._refs_path.exists():
+            try:
+                loaded = json.loads(self._refs_path.read_text(encoding="utf-8"))
+                if isinstance(loaded, dict):
+                    main_data = loaded
+            except Exception as e:
+                logger.warning("Failed to load MSTeams conversation refs: {}", e)
+
+        if meta_exists:
+            try:
+                loaded_meta = json.loads(self._refs_meta_path.read_text(encoding="utf-8"))
+                if isinstance(loaded_meta, dict):
+                    meta_data = loaded_meta
+            except Exception as e:
+                logger.warning("Failed to load MSTeams conversation refs metadata: {}", e)
+
+        return main_data, meta_data, meta_exists
+
+    def _load_refs_from_disk(self) -> dict[str, ConversationRef]:
+        """Load refs from disk with compatibility fallback for legacy layouts."""
+        main_data, meta_data, meta_exists = self._load_refs_raw()
+        if not main_data:
+            return {}
+
+        out: dict[str, ConversationRef] = {}
+        now = time.time()
+        for key, value in main_data.items():
+            ref = self._normalize_ref_record(value)
+            if not ref:
+                continue
+
+            meta_entry = meta_data.get(key) if isinstance(meta_data, dict) else None
+            meta_ts = None
+            if isinstance(meta_entry, dict):
+                meta_ts = self._safe_float(meta_entry.get("updated_at"))
+            elif meta_entry is not None:
+                meta_ts = self._safe_float(meta_entry)
+
+            if meta_ts is not None:
+                ref.updated_at = meta_ts
+            elif not meta_exists:
+                # First run after introducing meta sidecar: keep legacy refs alive
+                # by initializing timestamps to "now" instead of purging immediately.
+                ref.updated_at = now
+            elif ref.updated_at is None:
+                ref.updated_at = now
+
+            out[key] = ref
+        return out
+
     def _load_refs(self) -> dict[str, ConversationRef]:
         """Load stored conversation references."""
-        if not self._refs_path.exists():
-            return {}
-        try:
-            data = json.loads(self._refs_path.read_text(encoding="utf-8"))
-            out: dict[str, ConversationRef] = {}
-            for key, value in data.items():
-                out[key] = ConversationRef(**value)
-            return out
-        except Exception as e:
-            logger.warning("Failed to load MSTeams conversation refs: {}", e)
-            return {}
+        return self._load_refs_from_disk()
 
-    def _save_refs(self) -> None:
-        """Persist conversation references."""
+    @contextmanager
+    def _refs_file_lock(self):
+        """Cross-process lock while merging and writing refs state."""
+        self._refs_path.parent.mkdir(parents=True, exist_ok=True)
+        lock_fp = self._refs_lock_path.open("a+", encoding="utf-8")
         try:
-            stale_keys = [
-                key
-                for key, ref in self._conversation_refs.items()
-                if self._is_stale_or_unsupported_ref(ref)
-            ]
-            for key in stale_keys:
-                self._conversation_refs.pop(key, None)
+            if fcntl is not None:
+                fcntl.flock(lock_fp.fileno(), fcntl.LOCK_EX)
+            yield
+        finally:
+            try:
+                if fcntl is not None:
+                    fcntl.flock(lock_fp.fileno(), fcntl.LOCK_UN)
+            finally:
+                lock_fp.close()
 
-            data = {
-                key: {
-                    "service_url": ref.service_url,
-                    "conversation_id": ref.conversation_id,
-                    "bot_id": ref.bot_id,
-                    "activity_id": ref.activity_id,
-                    "conversation_type": ref.conversation_type,
-                    "tenant_id": ref.tenant_id,
-                    "updated_at": ref.updated_at,
+    def _is_webchat_service_url(self, service_url: str) -> bool:
+        """Return True when service URL points to unsupported Bot Framework Web Chat."""
+        normalized = service_url.strip()
+        if not normalized:
+            return False
+        host = (urlparse(normalized).hostname or "").strip().lower()
+        if host:
+            return host == MSTEAMS_WEBCHAT_HOST or host.endswith(f".{MSTEAMS_WEBCHAT_HOST}")
+        return MSTEAMS_WEBCHAT_HOST in normalized.lower()
+
+    def _prune_conversation_refs(self, *, now: float | None = None) -> bool:
+        """Remove stale and unsupported conversation refs from memory."""
+        if not self._conversation_refs:
+            return False
+
+        now_ts = time.time() if now is None else now
+        ttl_days = int(self.config.ref_ttl_days)
+        stale_before = now_ts - (ttl_days * 24 * 60 * 60)
+        keys_to_drop: list[str] = []
+
+        for key, ref in self._conversation_refs.items():
+            if self.config.prune_web_chat_refs and self._is_webchat_service_url(ref.service_url):
+                keys_to_drop.append(key)
+                continue
+
+            conv_type = str(ref.conversation_type or "").strip().lower()
+            if self.config.prune_non_personal_refs and conv_type and conv_type != "personal":
+                keys_to_drop.append(key)
+                continue
+
+            try:
+                updated_at = float(ref.updated_at) if ref.updated_at is not None else 0.0
+            except (TypeError, ValueError):
+                updated_at = 0.0
+            if updated_at <= 0 or updated_at < stale_before:
+                keys_to_drop.append(key)
+
+        if not keys_to_drop:
+            return False
+
+        for key in keys_to_drop:
+            self._conversation_refs.pop(key, None)
+        logger.info(
+            "MSTeams pruned {} stale/unsupported conversation refs (ttl={} days)",
+            len(keys_to_drop),
+            ttl_days,
+        )
+        return True
+
+    def _merge_refs_from_disk_locked(self) -> None:
+        """Merge disk refs into memory to reduce lost updates across processes."""
+        disk_refs = self._load_refs_from_disk()
+        for key, disk_ref in disk_refs.items():
+            mem_ref = self._conversation_refs.get(key)
+            if mem_ref is None:
+                self._conversation_refs[key] = disk_ref
+                continue
+            disk_ts = self._safe_float(disk_ref.updated_at) or 0.0
+            mem_ts = self._safe_float(mem_ref.updated_at) or 0.0
+            if disk_ts > mem_ts:
+                self._conversation_refs[key] = disk_ref
+
+    def _touch_conversation_ref(self, chat_id: str, *, persist: bool = False) -> None:
+        """Refresh updated_at for an active ref to keep it from expiring while used."""
+        with self._refs_guard:
+            ref = self._conversation_refs.get(str(chat_id))
+            if not ref:
+                return
+            now = time.time()
+            prev = self._safe_float(ref.updated_at) or 0.0
+            min_interval = max(0, int(self.config.ref_touch_interval_s))
+            if min_interval > 0 and prev > 0 and now - prev < min_interval:
+                return
+            ref.updated_at = now
+            if persist:
+                self._save_refs_locked()
+
+    def _write_json_atomically(self, path, data: dict[str, Any]) -> None:
+        """Write refs JSON atomically to reduce corruption risk during crashes."""
+        payload = json.dumps(data, indent=2)
+        tmp_path: str | None = None
+        try:
+            fd, tmp_path = tempfile.mkstemp(
+                dir=str(path.parent),
+                prefix=f"{path.name}.",
+                suffix=".tmp",
+            )
+            with os.fdopen(fd, "w", encoding="utf-8") as f:
+                f.write(payload)
+                f.flush()
+                os.fsync(f.fileno())
+            os.replace(tmp_path, path)
+        finally:
+            if tmp_path and os.path.exists(tmp_path):
+                try:
+                    os.unlink(tmp_path)
+                except OSError:
+                    pass
+
+    def _save_refs_locked(self, *, prune: bool = True) -> None:
+        """Persist conversation references (caller must hold _refs_guard)."""
+        try:
+            with self._refs_file_lock():
+                self._merge_refs_from_disk_locked()
+                if prune:
+                    self._prune_conversation_refs()
+                refs_data = {
+                    key: {
+                        "service_url": ref.service_url,
+                        "conversation_id": ref.conversation_id,
+                        "bot_id": ref.bot_id,
+                        "activity_id": ref.activity_id,
+                        "conversation_type": ref.conversation_type,
+                        "tenant_id": ref.tenant_id,
+                    }
+                    for key, ref in self._conversation_refs.items()
                 }
-                for key, ref in self._conversation_refs.items()
-            }
-            self._refs_path.write_text(json.dumps(data, indent=2), encoding="utf-8")
+                refs_meta = {
+                    key: {
+                        "updated_at": self._safe_float(ref.updated_at),
+                    }
+                    for key, ref in self._conversation_refs.items()
+                }
+                self._write_json_atomically(self._refs_path, refs_data)
+                self._write_json_atomically(self._refs_meta_path, refs_meta)
         except Exception as e:
             logger.warning("Failed to save MSTeams conversation refs: {}", e)
 
-    def _is_stale_or_unsupported_ref(self, ref: ConversationRef) -> bool:
-        """Reject unsupported refs and prune old refs."""
-        service_url = (ref.service_url or "").strip().lower()
-        conversation_type = (ref.conversation_type or "").strip().lower()
-        updated_at = ref.updated_at or 0.0
-        max_age_seconds = 30 * 24 * 60 * 60
-
-        if "webchat.botframework.com" in service_url:
-            return True
-        if conversation_type and conversation_type != "personal":
-            return True
-        if updated_at and updated_at < time.time() - max_age_seconds:
-            return True
-        return False
+    def _save_refs(self, *, prune: bool = True) -> None:
+        """Persist conversation references."""
+        with self._refs_guard:
+            self._save_refs_locked(prune=prune)
 
     async def _get_access_token(self) -> str:
         """Fetch an access token for Bot Framework / Azure Bot auth."""