hotfix: OAuth token caching by user identity instead of by MCPApp session ID (lastmile-ai#570)

* Fix workflow resume issue

* temp wip

* Fix workflow resume issue

* temp wip

* add audience validation

* add comfigured token support; add workflow_pre_auth; oauth example against github

* working oauth example with workflow_pre_auth

* fixes to oauth discovery; add dynamic oath example

* full e2e workflow

* improve how we're dealing with no oauth user

* all tests passing

* reformat

* rework oauth flow

* Update readme

* Fix user conflation on temporal; start of moving cache to app side

* cache tokens requested from temporal flow in app

* Implement local loopback OAuth callback server for MCPApp client-only runs

* Tests and more fixes for loopback, including browser launch

* various fixes

* additional fixes suggested by cursor

* merge and format

* Various updates to get the "interactive_tool" example working

* All examples working

* Regenerate schema and add docstrings

* Remove TESTING_GUIDE.md

* Fixes to make sure oauth identities are properly isolated across multiple users

* remove

* update org name

---------

Co-authored-by: Roman van der Krogt <[email protected]>

Author: 2 people

examples/oauth/interactive_tool/client.py

@@ -86,7 +86,7 @@ async def main() -> None:
             print("[client] Invoking github_org_search...")
             result = await connection.call_tool(
                 "github_org_search",
-                {"query": "lastmileai"},
+                {"query": "lastmile-ai"},
             )
             print("[client] Result:")
             for item in result.content or []:

examples/oauth/interactive_tool/server.py

@@ -59,7 +59,7 @@
 
 settings = Settings(
     execution_engine="asyncio",
-    logger=LoggerSettings(level="info"),
+    logger=LoggerSettings(level="debug"),
     oauth=OAuthSettings(
         callback_base_url=AnyHttpUrl("http://localhost:8000"),
         flow_timeout_seconds=300,

src/mcp_agent/core/context.py

@@ -4,10 +4,10 @@
 
 import asyncio
 import concurrent.futures
-from typing import Any, List, Optional, TYPE_CHECKING, Literal
+from typing import Any, Dict, List, Optional, TYPE_CHECKING, Literal
 import warnings
 
-from pydantic import ConfigDict
+from pydantic import ConfigDict, Field
 
 from mcp import ServerSession
 from mcp.server.fastmcp import FastMCP
@@ -34,6 +34,7 @@
 from mcp_agent.workflows.llm.llm_selector import ModelSelector
 from mcp_agent.logging.logger import get_logger
 from mcp_agent.tracing.token_counter import TokenCounter
+from mcp_agent.oauth.identity import OAuthUserIdentity
 
 
 if TYPE_CHECKING:
@@ -103,6 +104,7 @@ class Context(MCPContext):
     # OAuth helpers for downstream servers
     token_store: Optional[TokenStore] = None
     token_manager: Optional[TokenManager] = None
+    identity_registry: Dict[str, OAuthUserIdentity] = Field(default_factory=dict)
 
     model_config = ConfigDict(
         extra="allow",

src/mcp_agent/logging/logger.py

@@ -370,6 +370,7 @@ class LoggingConfig:
 
     _initialized: bool = False
     _event_filter_ref: EventFilter | None = None
+    _upstream_event_filter_ref: EventFilter | None = None
 
     @classmethod
     async def configure(
@@ -392,8 +393,11 @@ async def configure(
         """
         bus = AsyncEventBus.get(transport=transport)
         # Keep a reference to the provided filter so we can update at runtime
-        if event_filter is not None:
-            cls._event_filter_ref = event_filter
+        if event_filter is None:
+            event_filter = EventFilter()
+
+        cls._event_filter_ref = event_filter
+        cls._upstream_event_filter_ref = event_filter.model_copy(deep=True)
 
         # If already initialized, ensure critical listeners exist and return
         if cls._initialized:
@@ -411,7 +415,9 @@ async def configure(
                     MCP_UPSTREAM_LISTENER_NAME: _Final[str] = "mcp_upstream"
                     bus.add_listener(
                         MCP_UPSTREAM_LISTENER_NAME,
-                        MCPUpstreamLoggingListener(event_filter=cls._event_filter_ref),
+                        MCPUpstreamLoggingListener(
+                            event_filter=cls._upstream_event_filter_ref
+                        ),
                     )
             except Exception:
                 pass
@@ -451,7 +457,9 @@ async def configure(
                 MCP_UPSTREAM_LISTENER_NAME: Final[str] = "mcp_upstream"
                 bus.add_listener(
                     MCP_UPSTREAM_LISTENER_NAME,
-                    MCPUpstreamLoggingListener(event_filter=event_filter),
+                    MCPUpstreamLoggingListener(
+                        event_filter=cls._upstream_event_filter_ref
+                    ),
                 )
         except Exception:
             # Non-fatal if import fails
@@ -472,7 +480,7 @@ async def shutdown(cls):
     @classmethod
     def set_min_level(cls, level: EventType | str) -> None:
         """Update the minimum logging level on the shared event filter, if available."""
-        if cls._event_filter_ref is None:
+        if cls._upstream_event_filter_ref is None:
             return
         # Normalize level
         normalized = str(level).lower()
@@ -488,7 +496,7 @@ def set_min_level(cls, level: EventType | str) -> None:
             "alert": "error",
             "emergency": "error",
         }
-        cls._event_filter_ref.min_level = mapping.get(normalized, "info")
+        cls._upstream_event_filter_ref.min_level = mapping.get(normalized, "info")
 
     @classmethod
     def get_event_filter(cls) -> EventFilter | None:

src/mcp_agent/oauth/manager.py

@@ -342,6 +342,13 @@ async def get_access_token_if_present(
             self._default_identity,
         ]
         identities = _dedupe(identity_candidates)
+        logger.debug(
+            "Resolved identity candidates for token acquisition",
+            data={
+                "server": server_name,
+                "candidates": [candidate.cache_key for candidate in identities],
+            },
+        )
         if not identities:
             raise MissingUserIdentityError(
                 "No authenticated user available for OAuth authorization"
@@ -364,6 +371,14 @@ async def get_access_token_if_present(
             async with lock:
                 record = await self._token_store.get(key)
                 if record and not record.is_expired(leeway_seconds=leeway):
+                    logger.debug(
+                        "Token cache hit",
+                        data={
+                            "server": server_name,
+                            "identity": identity.cache_key,
+                            "resource": resolved.resource,
+                        },
+                    )
                     return record
 
                 if record and record.refresh_token:
@@ -546,6 +561,14 @@ async def ensure_access_token(
             )
 
             await self._token_store.set(user_key, record)
+            logger.debug(
+                "Stored new access token via authorization flow",
+                data={
+                    "server": server_name,
+                    "identity": flow_identity.cache_key,
+                    "resource": resolved.resource,
+                },
+            )
             return record
 
     async def invalidate(
@@ -780,25 +803,56 @@ def _session_identity(self, context: "Context") -> OAuthUserIdentity | None:
         except Exception:
             in_temporal = False
 
-        session_id = None
+        # Temporal workflows/activities carry their own execution identity.
         if in_temporal:
-            # Base the identity on the Temporal workflow execution ID
             try:
                 from mcp_agent.executor.temporal.temporal_context import (
                     get_execution_id as _get_exec_id,
                 )
+                from mcp_agent.server import app_server
 
-                session_id = _get_exec_id()
+                execution_id = _get_exec_id()
+                if execution_id:
+                    identity = app_server._get_identity_for_execution(execution_id)
+                    if identity is not None:
+                        return identity
             except Exception:
                 pass
 
-        if not session_id:
+        session_id = None
+        if in_temporal:
             session_id = getattr(context, "session_id", None)
-        if not session_id:
-            app = getattr(context, "app", None)
-            if app is not None:
-                session_id = getattr(app, "_session_id_override", None)
 
-        if not session_id:
-            return None
-        return OAuthUserIdentity(provider="mcp-session", subject=str(session_id))
+        if session_id:
+            try:
+                from mcp_agent.server import app_server
+
+                identity = app_server.get_identity_for_session(session_id, context)
+                if identity is not None:
+                    logger.debug(
+                        "Resolved session identity from registry",
+                        data={
+                            "session_id": session_id,
+                            "identity": identity.cache_key,
+                        },
+                    )
+                    return identity
+            except Exception as exc:
+                logger.debug(
+                    "Failed to resolve session identity from registry",
+                    data={"session_id": session_id, "error": repr(exc)},
+                )
+            fallback = OAuthUserIdentity(
+                provider="mcp-session", subject=str(session_id)
+            )
+            logger.debug(
+                "Falling back to synthetic session identity",
+                data={"session_id": session_id, "identity": fallback.cache_key},
+            )
+            return fallback
+
+        logger.debug(
+            "TokenManager no session identity resolved",
+            data={"context_session_id": getattr(context, "session_id", None)},
+        )
+        return None

src/mcp_agent/server/app_server.py

@@ -147,24 +147,54 @@ def _resolve_identity_for_request(
     identity = _CURRENT_IDENTITY.get()
     if identity is None and execution_id:
         identity = _get_identity_for_execution(execution_id)
-    if identity is None and app_context is not None:
-        session_id = getattr(app_context, "session_id", None)
-        identity = _session_identity_from_value(session_id)
-    if identity is None and ctx is not None:
-        session_id = _extract_session_id_from_context(ctx)
-        identity = _session_identity_from_value(session_id)
-    if identity is None and app_context is None and ctx is not None:
+    request_session_id: str | None = None
+    if ctx is not None:
+        request_session_id = _extract_session_id_from_context(ctx)
+    if app_context is None and ctx is not None:
         app = _get_attached_app(ctx.fastmcp)
         if app is not None and getattr(app, "context", None) is not None:
             app_context = app.context
+    if identity is None and request_session_id:
+        resolved = get_identity_for_session(request_session_id, app_context)
+        if resolved:
+            logger.debug(
+                "Resolved identity from session registry",
+                data={
+                    "session_id": request_session_id,
+                    "identity": resolved.cache_key,
+                },
+            )
+            identity = resolved
     if identity is None and app_context is not None:
         session_id = getattr(app_context, "session_id", None)
-        identity = _session_identity_from_value(session_id)
+        if session_id and session_id != request_session_id:
+            identity = get_identity_for_session(session_id, app_context)
     if identity is None:
         identity = DEFAULT_PRECONFIGURED_IDENTITY
     return identity
 
 
+def get_identity_for_session(
+    session_id: str | None, app_context: "Context" | None = None
+) -> OAuthUserIdentity | None:
+    """Lookup the cached identity for a given MCP session."""
+    if not session_id:
+        return None
+    if app_context is not None:
+        try:
+            identity = app_context.identity_registry.get(session_id)
+            if identity is not None:
+                return identity
+        except Exception:
+            pass
+    else:
+        logger.debug(
+            "No app context provided when resolving session identity",
+            data={"session_id": session_id},
+        )
+    return _session_identity_from_value(session_id)
+
+
 class ServerContext(ContextDependent):
     """Context object for the MCP App server."""
 
@@ -250,14 +280,11 @@ def _set_upstream_from_request_ctx_if_available(ctx: MCPContext) -> None:
     # First, try to use the session property from the FastMCP Context
     session = None
     try:
-        session = (
-            ctx.session
-        )  # This accesses the property which returns ctx.request_context.session
+        session = ctx.session
     except (AttributeError, ValueError):
         # ctx.session property might raise ValueError if context not available
         pass
 
-    # Capture authenticated user information if available
     session_id = _extract_session_id_from_context(ctx)
     identity: OAuthUserIdentity | None = None
     try:
@@ -268,7 +295,6 @@ def _set_upstream_from_request_ctx_if_available(ctx: MCPContext) -> None:
     if isinstance(auth_user, AuthenticatedUser):
         access_token = getattr(auth_user, "access_token", None)
         if access_token is not None:
-            # Prefer enriched token instances but fall back to raw data if necessary
             try:
                 from mcp_agent.oauth.access_token import MCPAccessToken
 
@@ -277,38 +303,35 @@ def _set_upstream_from_request_ctx_if_available(ctx: MCPContext) -> None:
                 else:
                     token_dict = getattr(access_token, "model_dump", None)
                     if callable(token_dict):
-                        maybe_token = MCPAccessToken.model_validate(
-                            access_token.model_dump()
-                        )
-                identity = OAuthUserIdentity.from_access_token(maybe_token)
+                        maybe_token = MCPAccessToken.model_validate(token_dict())
+                        if maybe_token is not None:
+                            identity = OAuthUserIdentity.from_access_token(maybe_token)
             except Exception:
                 identity = None
 
+    app_context: "Context" | None = None
     app: MCPApp | None = _get_attached_app(ctx.fastmcp)
     if app is not None and getattr(app, "context", None) is not None:
+        app_context = app.context
         if session is not None:
-            app.context.upstream_session = session
-        if session_id and not getattr(app.context, "session_id", None):
-            app.context.session_id = session_id
-
-        app_session_id = getattr(app.context, "session_id", None)
-    else:
-        app_session_id = None
-
-    if app_session_id:
-        app_identity = _session_identity_from_value(app_session_id)
-        if identity is None or (
-            isinstance(identity, OAuthUserIdentity)
-            and identity.provider == "mcp-session"
-        ):
-            identity = app_identity
+            app_context.upstream_session = session
 
-    if identity is None:
+    if identity is None and session_id:
         identity = _session_identity_from_value(session_id)
 
     if identity is None:
         identity = DEFAULT_PRECONFIGURED_IDENTITY
 
+    if app_context is not None and session_id and identity is not None:
+        try:
+            app_context.identity_registry[session_id] = identity
+            logger.debug(
+                "Registered identity for session",
+                data={"session_id": session_id, "identity": identity.cache_key},
+            )
+        except Exception:
+            pass
+
     _set_current_identity(identity)