Make Context bound to the request, allowing shared MCPApp connections (#573)

* 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

* Make Context bound to the request, allowing shared MCPApp connections

* Fix failing tests

* Address PR feedback

* address more PR feedback, plus an example

* Move examples

* more fixes

---------

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

Author: saqadri

examples/mcp_agent_server/context_isolation/README.md

@@ -0,0 +1,65 @@
+# Context Isolation Demo
+
+This example shows how per-request context scoping prevents logs and
+notifications from bleeding between concurrent MCP clients.
+
+## Setup
+
+- Install the example dependencies from this folder:
+  ```bash
+  uv pip install -r examples/mcp_agent_server/context_isolation/requirements.txt
+  ```
+- Optional: adjust `mcp_agent.config.yaml` if you want to tweak logging transports or
+  register additional MCP backends.
+
+## Running the example
+
+1. Start the SSE server in one terminal:
+
+   ```bash
+   uv run python examples/mcp_agent_server/context_isolation/server.py
+   ```
+
+   The server listens on `http://127.0.0.1:8000/sse` and exposes a single tool
+   (`emit_log`) that logs messages using the request-scoped context.
+
+2. In a second terminal, run the clients script. It launches two concurrent
+   clients that connect to the server, set independent logging levels, and call
+   the tool.
+
+   ```bash
+   uv run python examples/mcp_agent_server/context_isolation/clients.py
+   ```
+
+   Each client prints the logs and `demo/echo` notifications it receives. Client
+   A (set to `debug`) sees all messages it emits, while client B (set to
+   `error`) only receives error-level output. Notifications are tagged with the
+   originating session so you can observe the strict separation between the two
+   clients.
+
+## Expected output
+
+- Server console highlights two `SetLevelRequest` operations (one per client) followed
+  by a pair of `CallToolRequest` entries. You should also see an `emit_log` workflow
+  execution for each client with parameters matching the client payloads.
+
+- Client A prints both `debug` and `info` log notifications (one per tool call) and
+  the `demo/echo` notification containing its session id:
+
+  ```text
+  [A] log debug: ...
+  [A] log info: Workflow emit_log started execution ...
+  [A] tool result: ... "level": "debug"
+  ```
+
+- Client B only prints the `error` log notification—even after the second tool call—
+  confirming that the per-session
+  log level (`error`) filters out the info/debug output:
+
+  ```text
+  [B] log error: ...
+  [B] tool result: ... "level": "error"
+  ```
+
+If Client B ever receives an `info` or `debug` log entry, the request-scoped logging
+override is not working and should be investigated.

examples/mcp_agent_server/context_isolation/clients.py

@@ -0,0 +1,104 @@
+"""Connect two clients concurrently to demonstrate context isolation."""
+
+from __future__ import annotations
+
+import asyncio
+from datetime import timedelta
+from typing import Any
+
+from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
+from mcp import ClientSession
+
+from mcp_agent.app import MCPApp
+from mcp_agent.config import MCPServerSettings, MCPSettings, Settings
+from mcp_agent.core.context import Context
+from mcp_agent.mcp.gen_client import gen_client
+from mcp_agent.mcp.mcp_agent_client_session import MCPAgentClientSession
+
+
+SERVER_NAME = "context-isolation-server"
+SERVER_URL = "http://127.0.0.1:8000/sse"
+
+
+async def run_client(
+    client_name: str,
+    log_level: str,
+    payloads: list[str],
+    *,
+    delay_between_calls: float = 0.5,
+) -> None:
+    """Connect to the server, set logging, and invoke the emit_log tool for each payload."""
+
+    settings = Settings(
+        execution_engine="asyncio",
+        mcp=MCPSettings(
+            servers={
+                SERVER_NAME: MCPServerSettings(
+                    name=SERVER_NAME,
+                    description="Context isolation demo server",
+                    transport="sse",
+                    url=SERVER_URL,
+                )
+            }
+        ),
+    )
+
+    app = MCPApp(name=f"client-{client_name}", settings=settings)
+
+    async with app.run() as running_app:
+        context = running_app.context
+
+        async def on_log(params: Any) -> None:
+            try:
+                message = params.data.get("message") if params.data else None
+            except Exception:
+                message = None
+            print(f"[{client_name}] log {params.level}: {message}")
+
+        class DemoClientSession(MCPAgentClientSession):
+            async def _received_notification(self, notification):  # type: ignore[override]
+                method = getattr(getattr(notification, "root", None), "method", None)
+                if method and method != "notifications/message":
+                    print(
+                        f"[{client_name}] notify {method}: {notification.model_dump()}"
+                    )
+                return await super()._received_notification(notification)
+
+        def make_session(
+            read_stream: MemoryObjectReceiveStream,
+            write_stream: MemoryObjectSendStream,
+            read_timeout_seconds: timedelta | None,
+            context: Context | None = None,
+        ) -> ClientSession:
+            return DemoClientSession(
+                read_stream=read_stream,
+                write_stream=write_stream,
+                read_timeout_seconds=read_timeout_seconds,
+                logging_callback=on_log,
+                context=context,
+            )
+
+        async with gen_client(
+            SERVER_NAME,
+            context.server_registry,
+            client_session_factory=make_session,
+        ) as server:
+            await server.set_logging_level(log_level)
+            for idx, payload in enumerate(payloads, start=1):
+                result = await server.call_tool(
+                    "emit_log",
+                    arguments={"level": log_level, "message": payload},
+                )
+                print(f"[{client_name}] call {idx} result: {result}")
+                await asyncio.sleep(delay_between_calls)
+
+
+async def main() -> None:
+    await asyncio.gather(
+        run_client("A", "debug", ["hello from A", "A second info"]),
+        run_client("B", "error", ["hello from B", "B second info"]),
+    )
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

examples/mcp_agent_server/context_isolation/mcp_agent.config.yaml

@@ -0,0 +1,11 @@
+$schema: ../../schema/mcp-agent.config.schema.json
+
+execution_engine: asyncio
+
+logger:
+  transports: [console]
+  level: info
+
+mcp:
+  servers: {}
+

examples/mcp_agent_server/context_isolation/requirements.txt

@@ -0,0 +1,5 @@
+# Core framework dependency
+mcp-agent @ file://../../../  # Link to the local mcp-agent project root
+
+# Additional helper packages used by the client script
+anyio

examples/mcp_agent_server/context_isolation/server.py

@@ -0,0 +1,40 @@
+"""Simple SSE server demonstrating per-client context isolation."""
+
+import asyncio
+
+from mcp_agent.app import MCPApp
+from mcp_agent.core.context import Context
+from mcp_agent.server.app_server import create_mcp_server_for_app
+
+
+app = MCPApp(name="context-isolation-server")
+
+
+@app.tool("emit_log")
+async def emit_log(context: Context, level: str = "info", message: str = "hi") -> dict:
+    """Log a message at the requested level and emit a notification."""
+
+    session = context.request_session_id or "unknown"
+    await context.log(level, f"[{session}] {message}")
+    try:
+        await context.send_notification(
+            "demo/echo",
+            {
+                "session": session,
+                "level": level,
+                "message": message,
+            },
+        )
+    except Exception:
+        pass
+    return {"logged": message, "level": level, "session": session}
+
+
+async def main() -> None:
+    async with app.run() as running_app:
+        server = create_mcp_server_for_app(running_app)
+        await server.run_sse_async()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

src/mcp_agent/core/context.py

@@ -35,6 +35,7 @@
 from mcp_agent.logging.logger import get_logger
 from mcp_agent.tracing.token_counter import TokenCounter
 from mcp_agent.oauth.identity import OAuthUserIdentity
+from mcp_agent.core.request_context import get_current_request_context
 
 
 if TYPE_CHECKING:
@@ -73,7 +74,6 @@ class Context(MCPContext):
     human_input_handler: Optional[HumanInputCallback] = None
     elicitation_handler: Optional[ElicitationCallback] = None
     signal_notification: Optional[SignalWaitCallback] = None
-    upstream_session: Optional[ServerSession] = None
     model_selector: Optional[ModelSelector] = None
     session_id: str | None = None
     app: Optional["MCPApp"] = None
@@ -105,12 +105,51 @@ class Context(MCPContext):
     token_store: Optional[TokenStore] = None
     token_manager: Optional[TokenManager] = None
     identity_registry: Dict[str, OAuthUserIdentity] = Field(default_factory=dict)
+    request_session_id: str | None = None
+    request_identity: OAuthUserIdentity | None = None
 
     model_config = ConfigDict(
         extra="allow",
         arbitrary_types_allowed=True,  # Tell Pydantic to defer type evaluation
     )
 
+    @property
+    def upstream_session(self) -> ServerSession | None:  # type: ignore[override]
+        """
+        Resolve the active upstream session, preferring the request-scoped clone.
+
+        The base application context keeps an optional session used by scripts or
+        tests that set MCPApp.upstream_session directly. During an MCP request the
+        request-bound context is stored in a ContextVar; whenever callers reach the
+        base context while that request is active we return the request's session
+        instead of whichever client touched the base context last.
+        """
+        request_ctx = get_current_request_context()
+        if request_ctx is not None:
+            if request_ctx is self:
+                return getattr(self, "_upstream_session", None)
+
+            current = request_ctx
+            while current is not None:
+                parent_ctx = getattr(current, "_parent_context", None)
+                if parent_ctx is self:
+                    return getattr(current, "_upstream_session", None)
+                current = parent_ctx
+
+        explicit = getattr(self, "_upstream_session", None)
+        if explicit is not None:
+            return explicit
+
+        parent = getattr(self, "_parent_context", None)
+        if parent is not None:
+            return getattr(parent, "_upstream_session", None)
+
+        return None
+
+    @upstream_session.setter
+    def upstream_session(self, value: ServerSession | None) -> None:
+        object.__setattr__(self, "_upstream_session", value)
+
     @property
     def mcp(self) -> FastMCP | None:
         return self.app.mcp if self.app else None
@@ -144,9 +183,10 @@ def session(self) -> ServerSession | None:
 
         Returns None when no session can be resolved (e.g., local scripts).
         """
-        # 1) Explicit upstream session set by app/workflow
-        if getattr(self, "upstream_session", None) is not None:
-            return self.upstream_session
+        # 1) Explicit upstream session set by app/workflow (handles request clones)
+        explicit = getattr(self, "upstream_session", None)
+        if explicit is not None:
+            return explicit
 
         # 2) Try request-scoped session from FastMCP Context (may raise outside requests)
         try:
@@ -216,6 +256,13 @@ def bind_request(
         """
         # Shallow copy to preserve references to registries/loggers while keeping isolation
         bound: Context = self.model_copy(deep=False)
+        object.__setattr__(bound, "_upstream_session", None)
+        try:
+            object.__setattr__(bound, "_parent_context", self)
+        except Exception:
+            pass
+        bound.request_session_id = None
+        bound.request_identity = None
         try:
             setattr(bound, "_request_context", request_context)
         except Exception:
@@ -260,11 +307,16 @@ async def log(
         """
         # If we have a live FastMCP request context, delegate to parent
         try:
-            # will raise if request_context is not available
             _ = self.request_context  # type: ignore[attr-defined]
-            return await super().log(level, message, logger_name=logger_name)  # type: ignore[misc]
         except Exception:
             pass
+        else:
+            try:
+                return await super().log(  # type: ignore[misc]
+                    level, message, logger_name=logger_name
+                )
+            except Exception:
+                pass
 
         # Fall back to local logger if available
         try:
@@ -305,12 +357,10 @@ async def read_resource(self, uri: Any) -> Any:  # type: ignore[override]
         """
         # Use the parent implementation if request-bound fastmcp is available
         try:
-            if getattr(self, "_fastmcp", None) is not None:
-                return await super().read_resource(uri)  # type: ignore[misc]
+            return await super().read_resource(uri)  # type: ignore[misc]
         except Exception:
             pass
 
-        # Fall back to app-managed FastMCP if present
         try:
             mcp = self.mcp
             if mcp is not None:
@@ -501,6 +551,9 @@ def get_current_context() -> Context:
     Synchronous initializer/getter for global application context.
     For async usage, use aget_current_context instead.
     """
+    request_ctx = get_current_request_context()
+    if request_ctx is not None:
+        return request_ctx
     global _global_context
     if _global_context is None:
         try: