I Built a Custom MCP Server on Databricks — Here’s What the Managed Version Can’t Do
Part 2 of the MCP on Databricks series
In Part 1, I covered what Genie MCP and DBSQL MCP give you out of the box — and where they stop.
The short version: both are stateless, the managed Genie MCP has no `conversation_id` control, and the moment you need multi-turn context or cross-session memory, you’ve hit a wall.
This post is about what I built to get past that wall.
The Problems I Set Out to Solve
After testing the managed Genie MCP in a real agentic workflow, I ran into four specific gaps that couldn’t be patched at the prompt level:
1. No conversation_id control in the per-space Genie MCP
Every query_space tool call hits a single cross-space endpoint (/api/2.0/mcp/genie/{space_id}) and starts a brand new Genie conversation under the hood — stateless by design. You lose per-space routing, fine-grained tool descriptions, and any sense of session continuity.
2. No cross-session memory
Even if LangGraph maintains state within a session via its checkpointer, once the session ends, the Genie-side context is gone. There’s nothing to resume from.
3. No custom routing logic
The managed MCP exposes all Genie spaces as tools but gives you zero control over how the agent picks between them. The tool descriptions are too generic to drive reliable routing.
4. STM lives outside MCP
Short-term memory is managed at the LangGraph layer, not at the MCP layer. This means the MCP server itself is stateless — each tool call is atomic and knows nothing about previous calls in the same session.
The managed Genie MCP can’t solve any of these. So I built a custom FastMCP server on top of the Genie REST API.
Architecture
Before diving into the code, here’s how the pieces fit together:
Press enter or click to view image in full size
The key insight: `conversation_id` is the bridge between LangGraph state and Genie’s conversation thread. Persist it, and you get continuity. Drop it, and you’re back to square one every call.
Building the Custom FastMCP Server
Step 1 — Basic Server Skeleton
from fastmcp import FastMCP
import requests
import os
import time
mcp = FastMCP("databricks-genie-mcp")
DATABRICKS_HOST = os.environ["DATABRICKS_HOST"]
DATABRICKS_TOKEN = os.environ["DATABRICKS_TOKEN"]
GENIE_SPACE_ID = os.environ["GENIE_SPACE_ID"]
HEADERS = {
"Authorization": f"Bearer {DATABRICKS_TOKEN}",
"Content-Type": "application/json"
}
def genie_api(endpoint: str, method: str = "GET", payload: dict = None):
url = f"{DATABRICKS_HOST}/api/2.0/genie/spaces/{GENIE_SPACE_ID}/{endpoint}"
response = requests.request(method, url, headers=HEADERS, json=payload)
response.raise_for_status()
return response.json()Nothing surprising here — FastMCP server wrapping the Genie REST API with a shared helper for authenticated requests.
Step 2 — The Core: conversation_id Persistence
This is the piece the managed Genie MCP simply doesn’t give you.
@mcp.tool()
def create_new_conversation() -> dict:
"""
Start a new Genie conversation and return the conversation_id.
Store this in LangGraph state to enable multi-turn context.
"""
result = genie_api(
"start-conversation",
method="POST",
payload={"content": "Session initialized."}
)
return {
"conversation_id": result["conversation"]["conversation_id"],
"initial_message_id": result["message"]["message_id"],
"message": "New conversation created. Pass this conversation_id
to all subsequent query_genie calls."
}
@mcp.tool()
def query_genie_with_context(question: str, conversation_id: str) -> dict:
"""
Query Genie within an existing conversation thread.
Requires conversation_id from a previous create_new_conversation call.
"""
# Send message to existing conversation
message_result = genie_api(
f"conversations/{conversation_id}/messages",
method="POST",
payload={"content": question}
)
message_id = message_result["message_id"]Poll for result (Genie is async)
for _ in range(30):
status = genie_api(
f"conversations/{conversation_id}/messages/{message_id}"
)
if status["status"] == "COMPLETED":
return {
"answer": status.get("content", ""),
"conversation_id": conversation_id,
"message_id": message_id
}
elif status["status"] == "FAILED":
return {"error": "Genie query failed", "conversation_id": conversation_id}
time.sleep(2)
return {"error": "Timeout waiting for Genie response"}The critical difference from the managed MCP: `conversation_id` is passed explicitly on every call. Genie sees this as a continuation of the same conversation — it has full context of everything said before.
Step 3 — Cross-Session Memory via History Retrieval
This is the tool the managed Genie MCP has no equivalent for.
@mcp.tool()
def get_conversation_history(conversation_id: str) -> dict:
"""
Retrieve full message history from a previous Genie conversation.
Use this at the start of a new session to hydrate context.
"""
result = genie_api(f"conversations/{conversation_id}/messages")
messages = []
for msg in result.get("messages", []):
messages.append({
"role": msg.get("role"),
"content": msg.get("content"),
"timestamp": msg.get("created_at")
})
return {
"conversation_id": conversation_id,
"message_count": len(messages),
"history": messages
}At the start of a new session, the LangGraph agent calls `get_conversation_history()` with the persisted `conversation_id` from the previous session. It hydrates the context before asking anything new.
This is what makes the system behave like it has memory — even though MCP itself is stateless.
Wiring It Into LangGraph
STM vs LTM — Clarifying the Layers
LangGraph handles STM. The persisted `conversation_id` bridges LTM. The custom MCP server connects both to Genie.
The LangGraph Graph
from langgraph.graph import StateGraph, END
from langgraph.checkpoint.memory import MemorySaver
from langchain_openai import ChatOpenAI
from typing import TypedDict, Optional
class AgentState(TypedDict):
messages: list
conversation_id: Optional[str] # Genie conversation_id persisted here
question: str
answer: Optional[str]
llm = ChatOpenAI(
base_url=f"{host}/serving-endpoints",
api_key=user_token,
model=LLM_ENDPOINT,
temperature=0,
)
def initialize_session(state: AgentState) -> AgentState:
"""
At session start - either create new Genie conversation
or retrieve history from a previous one.
"""
if state.get("conversation_id"):
# Resume existing conversation - hydrate context
history = get_conversation_history(state["conversation_id"])
state["messages"].append({
"role": "system",
"content": f"Resuming previous conversation. History: {history}"
})
else:
# Brand new session - create fresh Genie conversation
new_conv = create_new_conversation()
state["conversation_id"] = new_conv["conversation_id"]
return state
def query_genie_node(state: AgentState) -> AgentState:
"""Send the user's question to Genie with conversation context."""
result = query_genie_with_context(
question=state["question"],
conversation_id=state["conversation_id"]
)
state["answer"] = result.get("answer", "No answer returned")
return state
# Build graph
workflow = StateGraph(AgentState)
workflow.add_node("initialize", initialize_session)
workflow.add_node("query_genie", query_genie_node)
workflow.set_entry_point("initialize")
workflow.add_edge("initialize", "query_genie")
workflow.add_edge("query_genie", END)
# Checkpointer handles STM across steps within the same thread_id
checkpointer = MemorySaver()
graph = workflow.compile(checkpointer=checkpointer)
# Run - same thread_id = same session context
config = {"configurable": {"thread_id": "supply-chain-session-001"}}
result = graph.invoke(
{"question": "What was our forecast accuracy last quarter?",
"conversation_id": "your-persisted-conversation-id-here"}, # from previous session
config=config
)
print(result["answer"])The `thread_id` is what LangGraph uses to resume checkpointed state within a session. The `conversation_id` is what Genie uses to maintain conversation thread continuity. Both work together — neither is sufficient alone.
MCP client wiring
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client
server_params = StdioServerParameters(command="python", args=["genie_mcp_server.py"])
async with stdio_client(server_params) as (read, write):
async with ClientSession(read, write) as session:
await session.initialize()
tools = await session.list_tools() # binds your 3 custom toolsWhat This Unlocks vs Managed MCP
The managed MCP wins on speed of setup. The custom server wins on everything an agentic system actually needs.
Honest Limitations
Building a custom MCP server doesn’t solve everything. Things that are still unsolved or require additional work:
1. Polling complexity
Genie is asynchronous. Every query needs a polling loop. Tune the interval and timeout wrong and you either miss results or hit rate limits.
2. conversation_id lifecycle management
Genie conversations don’t live forever. You need a strategy for when they expire — detect the failure, create a new conversation, and decide how much history to carry forward.
3. SQL consistency is still a Genie problem
The custom server routes better and maintains context — but Genie’s underlying tendency to generate inconsistent SQL on semantically similar questions is a product limitation, not something the MCP layer can fix.
4. Token cost of history hydration
Pulling full conversation history at session start adds tokens to every LLM call. For long conversations, this gets expensive quickly. You need a summarization strategy for LTM.
If you’re building something similar or hitting different walls with Genie MCP, drop a comment — happy to compare notes.
# Databricks # MCP #Genie #LangGraph #FastMCP #Agentic AI #Data Engineering
