fix(hosting): isolate OBS cache keys and bound refresh locks

- Key the cache by an (agent_id, tenant_id) tuple instead of
  "agent_id:tenant_id", so IDs that contain a colon cannot share another
  identity's cached token. make_key (new in this PR) becomes private.
- Keep each refresh lock on its cache entry, so capacity eviction and
  invalidate_all() reclaim it. Eviction skips entries with a refresh in flight.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Copilot-Session: 5cbf5f6b-cc40-4b7e-a591-65848db73a12

Author: DheerajPannala

libraries/microsoft-agents-a365-observability-hosting/microsoft_agents_a365/observability/hosting/token_cache_helpers/agent_token_cache.py

@@ -11,7 +11,7 @@
 import logging
 import time
 from collections.abc import Awaitable, Callable, Sequence
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from inspect import isawaitable
 from threading import Lock
 
@@ -59,6 +59,7 @@ class _Entry:
         token: str | None = None
         expires_on_ms: float | None = None
         acquired_on_ms: float | None = None
+        lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False, compare=False)
 
     _default_refresh_skew_ms = 60_000
     _default_max_token_age_ms = 3_600_000
@@ -67,18 +68,17 @@ class _Entry:
 
     def __init__(self, observability_scopes: Sequence[str] | None = None) -> None:
         """Initialize the token cache."""
-        self._map: dict[str, AgenticTokenCache._Entry] = {}
-        self._key_locks: dict[str, asyncio.Lock] = {}
+        self._map: dict[tuple[str, str], AgenticTokenCache._Entry] = {}
         self._lock = Lock()
         self._observability_scopes = (
             None if observability_scopes is None else tuple(observability_scopes)
         )
         self._removed_registration_logged = False
 
     @staticmethod
-    def make_key(agent_id: str, tenant_id: str) -> str:
-        """Create a cache key for an agent and tenant."""
-        return f"{agent_id}:{tenant_id}"
+    def _make_key(agent_id: str, tenant_id: str) -> tuple[str, str]:
+        # A tuple key keeps identities apart even when an ID contains a separator character.
+        return (agent_id, tenant_id)
 
     def register_observability(
         self,
@@ -118,10 +118,9 @@ async def refresh_observability_token(
         if not agent_id or not agent_id.strip() or not tenant_id or not tenant_id.strip():
             raise ValueError("[AgenticTokenCache] Agent and tenant IDs are required")
 
-        key = self.make_key(agent_id, tenant_id)
-        lock = self._get_key_lock(key)
-        async with lock:
-            entry = self._get_or_create_entry(key)
+        key = self._make_key(agent_id, tenant_id)
+        entry = self._get_or_create_entry(key)
+        async with entry.lock:
             if entry.token is not None and not self._is_expired(entry):
                 return entry.token
 
@@ -133,7 +132,7 @@ async def get_observability_token(self, agent_id: str, tenant_id: str) -> str |
         This method is a pure cache read. It never acquires a token and never
         calls delegated token exchange.
         """
-        key = self.make_key(agent_id, tenant_id)
+        key = self._make_key(agent_id, tenant_id)
         with self._lock:
             entry = self._map.get(key)
 
@@ -147,7 +146,7 @@ async def get_observability_token(self, agent_id: str, tenant_id: str) -> str |
 
     def invalidate_token(self, agent_id: str, tenant_id: str) -> None:
         """Invalidate one cached token."""
-        key = self.make_key(agent_id, tenant_id)
+        key = self._make_key(agent_id, tenant_id)
         with self._lock:
             entry = self._map.get(key)
             if entry is not None:
@@ -158,7 +157,7 @@ def invalidate_all(self) -> None:
         with self._lock:
             self._map.clear()
 
-    def _get_or_create_entry(self, key: str) -> _Entry:
+    def _get_or_create_entry(self, key: tuple[str, str]) -> _Entry:
         with self._lock:
             entry = self._map.get(key)
             if entry is not None:
@@ -171,22 +170,22 @@ def _get_or_create_entry(self, key: str) -> _Entry:
                 raise ValueError("[AgenticTokenCache] No valid scopes")
 
             if len(self._map) >= self._max_cache_size:
-                oldest_key = next(iter(self._map), None)
-                if oldest_key is not None:
-                    del self._map[oldest_key]
+                # Evict the oldest idle entry; an entry with a refresh in flight keeps its lock.
+                idle_key = next(
+                    (
+                        existing_key
+                        for existing_key, existing in self._map.items()
+                        if not existing.lock.locked()
+                    ),
+                    None,
+                )
+                if idle_key is not None:
+                    del self._map[idle_key]
 
             entry = AgenticTokenCache._Entry(scopes=scopes)
             self._map[key] = entry
             return entry
 
-    def _get_key_lock(self, key: str) -> asyncio.Lock:
-        with self._lock:
-            lock = self._key_locks.get(key)
-            if lock is None:
-                lock = asyncio.Lock()
-                self._key_locks[key] = lock
-            return lock
-
     def _get_effective_scopes(self) -> tuple[str, ...]:
         scopes = self._observability_scopes
         if scopes is None:

tests/observability/hosting/token_cache_helpers/test_agent_token_cache.py

@@ -253,7 +253,7 @@ async def test_opaque_token_uses_fresh_fallback_ttl(token_cache):
     await token_cache.refresh_observability_token("agent", "tenant", lambda *_: "opaque-token")
 
     assert await token_cache.get_observability_token("agent", "tenant") == "opaque-token"
-    entry = token_cache._map[AgenticTokenCache.make_key("agent", "tenant")]
+    entry = token_cache._map[AgenticTokenCache._make_key("agent", "tenant")]
     entry.acquired_on_ms = (time.time() * 1000) - token_cache._default_max_token_age_ms - 1
     assert await token_cache.get_observability_token("agent", "tenant") is None
 
@@ -307,3 +307,61 @@ async def test_cache_evicts_oldest_entry_when_capacity_is_reached(token_cache):
     assert await token_cache.get_observability_token("one", "tenant") is None
     assert await token_cache.get_observability_token("two", "tenant") == token
     assert await token_cache.get_observability_token("three", "tenant") == token
+
+
+@pytest.mark.asyncio
+async def test_separator_bearing_ids_do_not_share_a_cache_entry(token_cache):
+    """IDs that contain a separator character never alias another identity."""
+    first = MagicMock(return_value="token-for-first")
+    second = MagicMock(return_value="token-for-second")
+
+    assert await token_cache.refresh_observability_token("a:b", "c", first) == "token-for-first"
+    assert await token_cache.refresh_observability_token("a", "b:c", second) == "token-for-second"
+
+    second.assert_called_once()
+    assert await token_cache.get_observability_token("a:b", "c") == "token-for-first"
+    assert await token_cache.get_observability_token("a", "b:c") == "token-for-second"
+
+
+@pytest.mark.asyncio
+async def test_identity_churn_keeps_every_per_identity_registry_bounded(token_cache):
+    """Refresh locks live and die with cache entries, so identity churn stays bounded."""
+    token_cache._max_cache_size = 2
+    token = make_jwt(300)
+    for agent_id in ("one", "two", "three", "four"):
+        await token_cache.refresh_observability_token(agent_id, "tenant", lambda *_: token)
+
+    registries = {
+        name: len(value) for name, value in vars(token_cache).items() if isinstance(value, dict)
+    }
+    assert registries == {"_map": 2}
+
+    token_cache.invalidate_all()
+    registries = {
+        name: len(value) for name, value in vars(token_cache).items() if isinstance(value, dict)
+    }
+    assert registries == {"_map": 0}
+
+
+@pytest.mark.asyncio
+async def test_eviction_keeps_entry_with_refresh_in_flight(token_cache):
+    """Capacity eviction skips an identity whose refresh is still running."""
+    token_cache._max_cache_size = 1
+    release = asyncio.Event()
+
+    async def slow_resolver(agent_id: str, tenant_id: str, scopes: list[str]) -> str:
+        await release.wait()
+        return "token-for-slow"
+
+    in_flight = asyncio.create_task(
+        token_cache.refresh_observability_token("slow", "tenant", slow_resolver)
+    )
+    await asyncio.sleep(0)
+    fast = await token_cache.refresh_observability_token(
+        "fast", "tenant", lambda *_: "token-for-fast"
+    )
+
+    release.set()
+    assert fast == "token-for-fast"
+    assert await in_flight == "token-for-slow"
+    assert await token_cache.get_observability_token("slow", "tenant") == "token-for-slow"