Add session creation and browsing to Mason chat UI

integrations/mason/src/databricks_mason/templates/mcp_runtime_langgraph.py

@@ -10,8 +10,11 @@
     downscope_wire,
     load_tools,
 )
+from agent.mason.workspace import (  # ty: ignore[unresolved-import]
+    workspace_client,
+    workspace_headers,
+)
 from agent.mcps import build_mcp_servers  # ty: ignore[unresolved-import]
-from databricks.sdk import WorkspaceClient
 from databricks_langchain import (  # ty: ignore[unresolved-import]
     DatabricksMCPServer,
     DatabricksMultiServerMCPClient,
@@ -22,13 +25,14 @@
 
 
 def _server_from_tool(tool: ToolRecord) -> DatabricksMCPServer | None:
-    workspace_client = WorkspaceClient()
-    host = workspace_client.config.host.rstrip("/")
+    client = workspace_client()
+    host = client.config.host.rstrip("/")
     if tool.kind in {"sandbox", "mcp"}:
         return DatabricksMCPServer(
             name=tool.id,
             url=f"{host}/ai-gateway/mcp-services/{tool.service}",
-            workspace_client=workspace_client,
+            headers=workspace_headers() or None,
+            workspace_client=client,
             timeout=120.0,
         )
     if tool.kind == "uc_function":
@@ -38,7 +42,8 @@ def _server_from_tool(tool: ToolRecord) -> DatabricksMCPServer | None:
             schema=schema,
             function_name=function_name,
             name=tool.id,
-            workspace_client=workspace_client,
+            headers=workspace_headers() or None,
+            workspace_client=client,
             timeout=120.0,
         )
     return None

integrations/mason/templates/agent-langgraph/README.md

@@ -146,9 +146,16 @@ poll reaches the same replica, but the run itself does not survive a restart.
 
 When initialized with `mason init --framework langgraph --enable-chat-app`, the browser also calls:
 
+- `POST /api/session/new` to generate a fresh session id and replace the routing cookie. The request
+  has no body-level `session_id`; the response includes the new and previous ids.
 - `POST /api/demo/sessions` to create or resolve the current cookie-backed managed session.
+- `GET /api/demo/sessions` to list recent sessions for the configured actor. In local in-memory mode
+  it returns only the current browser session.
+- `POST /api/demo/sessions/{session_id}/open` to verify an actor-scoped managed session, replace the
+  routing cookie, and load that session's transcript and pending state.
 - `GET /api/demo/session/items` to load the current transcript. Without a managed Session Store it
-  reconstructs messages and pending interrupts from the in-process LangGraph checkpoint.
+  reconstructs messages and pending interrupts from the in-process LangGraph checkpoint. Managed
+  responses filter out checkpoint fragments and durability events before returning items to the UI.
 - `POST /api/demo/session/items` to mirror user, assistant, tool, and human-decision items into the
   managed Session Store.
 - `GET /api/demo/memory/entries`, `POST /api/demo/memory/entries`, and

integrations/mason/templates/agent-langgraph/agent/agent.py

@@ -3,7 +3,6 @@
 from collections.abc import AsyncGenerator, AsyncIterator
 from typing import Any
 
-from databricks.sdk import WorkspaceClient
 from databricks_langchain import ChatDatabricks
 from langchain.agents import create_agent
 from langchain.agents.middleware import HumanInTheLoopMiddleware
@@ -13,6 +12,7 @@
 from agent.mason import mcp_runtime, tracing
 from agent.mason.memory import memory_tools
 from agent.mason.session_store import checkpointer, thread_config
+from agent.mason.workspace import workspace_client, workspace_headers
 
 # Importing the tools package auto-registers every tool module.
 from agent.tools import all_tools
@@ -28,6 +28,16 @@
 REQUIRE_APPROVAL = {"send_message": True}
 
 
+class _RoutedChatDatabricks(ChatDatabricks):
+    """Forward account-host workspace routing to the underlying OpenAI clients."""
+
+    def _get_client_kwargs(self) -> dict[str, Any]:
+        kwargs = super()._get_client_kwargs()
+        if headers := workspace_headers():
+            kwargs["default_headers"] = headers
+        return kwargs
+
+
 def configure() -> None:
     """Wire up global state; call once at server startup (not at import)."""
     _check_databricks_auth()
@@ -42,7 +52,7 @@ def _check_databricks_auth() -> None:
     the model client uses, so the failure is immediate and actionable.
     """
     try:
-        WorkspaceClient()
+        workspace_client()
     except Exception as e:
         profile = os.getenv("DATABRICKS_CONFIG_PROFILE")
         target = (
@@ -65,7 +75,7 @@ async def create_agent_graph():
         [HumanInTheLoopMiddleware(interrupt_on=REQUIRE_APPROVAL)] if REQUIRE_APPROVAL else []
     )
     return create_agent(
-        model=ChatDatabricks(endpoint=MODEL),
+        model=_RoutedChatDatabricks(endpoint=MODEL, workspace_client=workspace_client()),
         tools=tools,
         middleware=middleware,
         checkpointer=checkpointer(),

integrations/mason/templates/agent-langgraph/agent/mason/durability.py

@@ -13,6 +13,8 @@
 
 from databricks.sdk import WorkspaceClient
 
+from agent.mason.workspace import workspace_client as default_workspace_client
+
 _EVENT_TYPE = "mason_demo_durability"
 _SESSION_STORE_ENV = "AGENT_SESSION_STORE"
 _SESSION_ACTOR_ENV = "AGENT_SESSION_ACTOR_ID"
@@ -35,7 +37,7 @@ class _SessionStoreClient:
     """Minimal Session Store client kept local to the generated durability demo."""
 
     def __init__(self, workspace_client: WorkspaceClient | None = None) -> None:
-        self._workspace = workspace_client or WorkspaceClient()
+        self._workspace = workspace_client or default_workspace_client()
         self._store_name = ""
 
     def set_session_store(self, session_store_name: str) -> "_SessionStoreClient":

integrations/mason/templates/agent-langgraph/agent/mason/mcp_runtime.py

@@ -5,24 +5,25 @@
 import logging
 from typing import Any
 
-from databricks.sdk import WorkspaceClient
 from databricks_langchain import DatabricksMCPServer, DatabricksMultiServerMCPClient
 from langchain_mcp_adapters.sessions import create_session
 
 from agent.mason.tool_manifest import ToolRecord, downscope_wire, load_tools
+from agent.mason.workspace import workspace_client, workspace_headers
 from agent.mcps import build_mcp_servers
 
 logger = logging.getLogger(__name__)
 
 
 def _server_from_tool(tool: ToolRecord) -> DatabricksMCPServer | None:
-    workspace_client = WorkspaceClient()
-    host = workspace_client.config.host.rstrip("/")
+    client = workspace_client()
+    host = client.config.host.rstrip("/")
     if tool.kind in {"sandbox", "mcp"}:
         return DatabricksMCPServer(
             name=tool.id,
             url=f"{host}/ai-gateway/mcp-services/{tool.service}",
-            workspace_client=workspace_client,
+            headers=workspace_headers() or None,
+            workspace_client=client,
             timeout=120.0,
         )
     if tool.kind == "uc_function":
@@ -32,7 +33,8 @@ def _server_from_tool(tool: ToolRecord) -> DatabricksMCPServer | None:
             schema=schema,
             function_name=function_name,
             name=tool.id,
-            workspace_client=workspace_client,
+            headers=workspace_headers() or None,
+            workspace_client=client,
             timeout=120.0,
         )
     return None

integrations/mason/templates/agent-langgraph/agent/mason/memory.py

@@ -15,9 +15,10 @@
 
 import os
 
-from databricks.sdk import WorkspaceClient
 from langchain_core.tools import BaseTool, tool
 
+from agent.mason.workspace import workspace_client
+
 _AGENTS_V1 = "/api/agents/v1"
 
 
@@ -31,7 +32,7 @@ def _store_path() -> str:
 
 def _api():
     # Build the client lazily (needs workspace auth) so importing this module stays cheap.
-    return WorkspaceClient().api_client
+    return workspace_client().api_client
 
 
 @tool

integrations/mason/templates/agent-langgraph/agent/mason/session_store_client.py

@@ -17,6 +17,8 @@
 
 from databricks.sdk import WorkspaceClient
 
+from agent.mason.workspace import workspace_client as default_workspace_client
+
 _API_ROOT = "/api/agents/v1"
 
 
@@ -42,7 +44,7 @@ class SessionStoreClient:
     """Thin REST client over the managed Session Store API."""
 
     def __init__(self, workspace_client: Optional[WorkspaceClient] = None) -> None:
-        self._api = (workspace_client or WorkspaceClient()).api_client
+        self._api = (workspace_client or default_workspace_client()).api_client
         self._store_name: Optional[str] = None
 
     def set_session_store(self, session_store_name: str) -> "SessionStoreClient":

integrations/mason/templates/agent-langgraph/agent/mason/workspace.py

@@ -0,0 +1,27 @@
+"""Construct Databricks SDK clients with workspace routing when required."""
+
+from __future__ import annotations
+
+import os
+
+from databricks.sdk import WorkspaceClient
+
+
+def workspace_headers() -> dict[str, str]:
+    """Return headers required to route an account-host request to its workspace."""
+    workspace_id = os.getenv("DATABRICKS_WORKSPACE_ID", "").strip()
+    return {"X-Databricks-Org-Id": workspace_id} if workspace_id else {}
+
+
+def workspace_client() -> WorkspaceClient:
+    """Return the environment-authenticated client for the active workspace.
+
+    ``databricks apps run-local`` can authenticate through an account-level vanity host while
+    exposing the target workspace through ``DATABRICKS_WORKSPACE_ID``. The SDK needs the same
+    routing header as the Databricks CLI for those profiles. Ordinary workspace hosts and deployed
+    Apps continue to use the SDK's default authentication chain.
+    """
+    headers = workspace_headers()
+    if not headers:
+        return WorkspaceClient()
+    return WorkspaceClient(custom_headers=headers)

integrations/mason/templates/agent-langgraph/runtime/runtime.py

@@ -15,9 +15,10 @@
 
 Endpoints: ``POST /invocations`` (``stream: true`` → SSE ending with ``data: [DONE]``;
 ``background: true`` → an ``invocation_id`` to poll), ``GET /invocations/{invocation_id}`` to poll a
-background run, and ``GET /health``. Each route also has an ``/api`` alias because Databricks Apps
-accepts programmatic Bearer-token authentication only on paths under ``/api/``. Each request is
-wrapped in an MLflow span for tracing.
+background run, ``POST /api/session/new`` to rotate the routing session, and ``GET /health``. The
+invocation and health routes also have ``/api`` aliases because Databricks Apps accepts programmatic
+Bearer-token authentication only on paths under ``/api/``. Each request is wrapped in an MLflow span
+for tracing.
 """
 
 import asyncio
@@ -26,7 +27,7 @@
 
 import mlflow
 from agent.mason.background import BackgroundRuns
-from fastapi import FastAPI, Request
+from fastapi import FastAPI, Request, Response
 from fastapi.responses import JSONResponse, StreamingResponse
 from uuid_utils import uuid7
 
@@ -55,6 +56,27 @@ def _set_trace_name(name: str) -> None:
         mlflow.update_current_trace(tags={_TRACE_NAME_TAG: name})
 
 
+def rotate_session_cookie(request: Request, response: Response, session_id: str) -> None:
+    if request.cookies.get(_ROUTING_COOKIE):
+        response.set_cookie(
+            _ROUTING_COOKIE,
+            session_id,
+            secure=True,
+            httponly=True,
+            samesite="lax",
+            path="/",
+        )
+        response.delete_cookie(_LOCAL_SESSION_COOKIE, path="/")
+    elif request.cookies.get(_LOCAL_SESSION_COOKIE):
+        response.set_cookie(
+            _LOCAL_SESSION_COOKIE,
+            session_id,
+            httponly=True,
+            samesite="lax",
+            path="/",
+        )
+
+
 def build_app(invoke_handler: InvokeHandler, stream_handler: StreamHandler) -> FastAPI:
     """Build the FastAPI app wiring the endpoints to the agent's invoke/stream handlers."""
     app = FastAPI(title="Agent Server")
@@ -136,4 +158,18 @@ async def retrieve(invocation_id: str):
     async def health() -> dict[str, str]:
         return {"status": "ok"}
 
+    @app.post("/api/session/new")
+    async def new_session(request: Request) -> JSONResponse:
+        previous_session_id = request.state.session_id
+        session_id = str(uuid7())
+        request.state.session_id = session_id
+        response = JSONResponse(
+            {
+                "session_id": session_id,
+                "previous_session_id": previous_session_id,
+            }
+        )
+        rotate_session_cookie(request, response, session_id)
+        return response
+
     return app