Implementing MCP with LangGraph: A Practical Walkthrough

from typing import TypedDict, List, Dict, Optional, Annotatedfrom dataclasses import dataclassfrom datetime import datetime@dataclassclass ContextReference:    """    A reference to context managed by MCP.    This is what goes in graph state.    """    key: str                    # Unique identifier    version: str                # Version hash    source: str                 # Which MCP server    expires_at: Optional[datetime] = None    metadata: Dict = None        def to_dict(self) -> Dict:        return {            "key": self.key,            "version": self.version,            "source": self.source,            "expires_at": self.expires_at.isoformat() if self.expires_at else None,            "metadata": self.metadata or {}        }class AgentState(TypedDict):    """    LangGraph state: small, serializable, reference-based.    """    # Workflow control    current_step: str    steps_completed: List[str]    decision_path: List[str]        # Context references (NOT actual data)    context_refs: Annotated[List[ContextReference], "Context references managed by MCP"]        # Execution metadata    user_id: str    session_id: str    started_at: datetime        # Results (small summaries, not full data)    summary: Optional[str]    final_response: Optional[str]        # Error tracking    errors: List[Dict]

from typing import Any, Dict, List, Optionalimport asyncioimport hashlibfrom datetime import datetime, timedeltaclass MCPContextManager:    """    Manages context lifecycle independently of graph execution.    This is your MCP integration layer.    """        def __init__(self, mcp_clients: Dict[str, Any], cache_ttl: int = 300):        self.mcp_clients = mcp_clients        self.cache_ttl = cache_ttl                # In-memory cache for hot context        # In production: Redis or similar        self._cache: Dict[str, tuple[Any, datetime]] = {}                # Version tracking        self._versions: Dict[str, str] = {}        async def resolve_context(        self,        ref: ContextReference,        use_cache: bool = True    ) -> Any:        """        Resolve a context reference to actual data.        This is called by graph nodes, not by the graph itself.        """        cache_key = f"{ref.source}:{ref.key}:{ref.version}"                # Check cache first        if use_cache and cache_key in self._cache:            data, cached_at = self._cache[cache_key]            if datetime.utcnow() - cached_at < timedelta(seconds=self.cache_ttl):                return data                # Cache miss: fetch from MCP server        client = self.mcp_clients.get(ref.source)        if not client:            raise ValueError(f"Unknown MCP source: {ref.source}")                data = await client.fetch_resource(ref.key)                # Cache the result        self._cache[cache_key] = (data, datetime.utcnow())                return data        async def create_context_reference(        self,        source: str,        key: str,        ttl_seconds: Optional[int] = None    ) -> ContextReference:        """        Create a new context reference with versioning.        Called when a node produces context that should be available to later nodes.        """        # Fetch current data to compute version        client = self.mcp_clients[source]        data = await client.fetch_resource(key)                # Version is hash of content        version = self._compute_version(data)                # Store version mapping        version_key = f"{source}:{key}"        self._versions[version_key] = version                expires_at = None        if ttl_seconds:            expires_at = datetime.utcnow() + timedelta(seconds=ttl_seconds)                return ContextReference(            key=key,            version=version,            source=source,            expires_at=expires_at,            metadata={"created_at": datetime.utcnow().isoformat()}        )        async def resolve_multiple(        self,        refs: List[ContextReference],        parallel: bool = True    ) -> Dict[str, Any]:        """        Resolve multiple context references efficiently.        Parallel fetching is the default—MCP shines here.        """        if not parallel:            results = {}            for ref in refs:                results[ref.key] = await self.resolve_context(ref)            return results                # Parallel resolution        tasks = {ref.key: self.resolve_context(ref) for ref in refs}        results = await asyncio.gather(*tasks.values(), return_exceptions=True)                # Handle failures gracefully        return {            key: result if not isinstance(result, Exception) else None            for key, result in zip(tasks.keys(), results)        }        def _compute_version(self, data: Any) -> str:        """        Compute version hash from data.        This ensures version changes when content changes.        """        # Serialize data deterministically        import json        serialized = json.dumps(data, sort_keys=True)        return hashlib.sha256(serialized.encode()).hexdigest()[:16]        async def refresh_if_stale(self, ref: ContextReference) -> ContextReference:        """        Check if context version has changed and return updated reference.        This is how you detect context updates.        """        version_key = f"{ref.source}:{ref.key}"        current_version = self._versions.get(version_key)                if current_version != ref.version:            # Version changed, create new reference            return await self.create_context_reference(                ref.source,                ref.key,                ttl_seconds=300            )                return ref

from langgraph.graph import StateGraph, ENDclass MCPAwareLangGraph:    """    LangGraph integration with MCP context management.    """        def __init__(self, context_manager: MCPContextManager):        self.context_manager = context_manager        self.graph = self._build_graph()        def _build_graph(self) -> StateGraph:        """Build the state graph with MCP-aware nodes"""        workflow = StateGraph(AgentState)                # Define nodes        workflow.add_node("analyze", self._analyze_node)        workflow.add_node("fetch_data", self._fetch_data_node)        workflow.add_node("process", self._process_node)        workflow.add_node("synthesize", self._synthesize_node)                # Define edges        workflow.set_entry_point("analyze")        workflow.add_edge("analyze", "fetch_data")        workflow.add_edge("fetch_data", "process")        workflow.add_edge("process", "synthesize")        workflow.add_edge("synthesize", END)                return workflow.compile()        async def _analyze_node(self, state: AgentState) -> AgentState:        """        Analyze query and determine what context is needed.        This node DOES NOT fetch data—it creates references.        """        # Determine required context sources        # This is fast business logic, not data fetching        required_sources = self._determine_required_sources(state)                # Create context references (doesn't fetch yet)        context_refs = []        for source, key in required_sources:            ref = await self.context_manager.create_context_reference(                source=source,                key=key,                ttl_seconds=300            )            context_refs.append(ref)                return {            **state,            "current_step": "analyze",            "steps_completed": state.get("steps_completed", []) + ["analyze"],            "context_refs": context_refs        }        async def _fetch_data_node(self, state: AgentState) -> AgentState:        """        Resolve context references to actual data.        This is where MCP actually fetches.        """        refs = state.get("context_refs", [])                # Parallel fetch through MCP        context_data = await self.context_manager.resolve_multiple(            refs,            parallel=True        )                # DON'T put data in state        # Instead, store it in MCP and update references                # For nodes that need the data, they'll resolve refs themselves        # State just tracks that fetching is complete                return {            **state,            "current_step": "fetch_data",            "steps_completed": state.get("steps_completed", []) + ["fetch_data"]        }        async def _process_node(self, state: AgentState) -> AgentState:        """        Process context data.        This node resolves references as needed.        """        refs = state.get("context_refs", [])                # Resolve only the context we need for this node        # Other nodes might need different context        needed_refs = [r for r in refs if r.source in ["database", "api"]]        context_data = await self.context_manager.resolve_multiple(needed_refs)                # Do actual processing        # This is where your business logic goes        processed_result = self._process_business_logic(context_data)                # Store large result back in MCP        result_ref = await self._store_result_in_mcp(processed_result)                # Update state with reference to result        updated_refs = state.get("context_refs", []) + [result_ref]                return {            **state,            "current_step": "process",            "steps_completed": state.get("steps_completed", []) + ["process"],            "context_refs": updated_refs        }        async def _synthesize_node(self, state: AgentState) -> AgentState:        """        Final synthesis using all available context.        """        refs = state.get("context_refs", [])                # Resolve all context for final synthesis        all_context = await self.context_manager.resolve_multiple(refs)                # Generate final response        # In production: call your LLM here with assembled context        final_response = self._generate_final_response(all_context, state)                return {            **state,            "current_step": "synthesize",            "steps_completed": state.get("steps_completed", []) + ["synthesize"],            "final_response": final_response        }        async def _store_result_in_mcp(self, result: Any) -> ContextReference:        """        Store large results in MCP instead of state.        """        # Generate unique key for this result        result_key = f"result_{datetime.utcnow().timestamp()}"                # Store in MCP (implementation depends on your MCP server)        # For now: add to cache        version = self.context_manager._compute_version(result)        cache_key = f"results:{result_key}:{version}"        self.context_manager._cache[cache_key] = (result, datetime.utcnow())                return ContextReference(            key=result_key,            version=version,            source="results",            metadata={"type": "processed_result"}        )

class ContextVersionManager:    """    Manages context versioning for consistency and freshness trade-offs.    """        def __init__(self, context_manager: MCPContextManager):        self.context_manager = context_manager        async def ensure_consistency(        self,        state: AgentState,        strategy: str = "snapshot"    ) -> AgentState:        """        Ensure context consistency across workflow execution.                Strategies:        - snapshot: Use versions from start, ignore updates        - latest: Always use latest versions        - hybrid: Use snapshot for critical context, latest for auxiliary        """        refs = state.get("context_refs", [])                if strategy == "snapshot":            # Keep original versions, don't refresh            return state                elif strategy == "latest":            # Refresh all context to latest versions            updated_refs = []            for ref in refs:                updated_ref = await self.context_manager.refresh_if_stale(ref)                updated_refs.append(updated_ref)                        return {**state, "context_refs": updated_refs}                elif strategy == "hybrid":            # Refresh only non-critical context            critical_sources = {"database", "api"}            updated_refs = []                        for ref in refs:                if ref.source in critical_sources:                    # Keep original version for consistency                    updated_refs.append(ref)                else:                    # Use latest for auxiliary context                    updated_ref = await self.context_manager.refresh_if_stale(ref)                    updated_refs.append(updated_ref)                        return {**state, "context_refs": updated_refs}                return state        async def detect_context_drift(        self,        old_state: AgentState,        new_state: AgentState    ) -> List[Dict]:        """        Detect which context has changed between states.        Useful for debugging and observability.        """        old_refs = {r.key: r for r in old_state.get("context_refs", [])}        new_refs = {r.key: r for r in new_state.get("context_refs", [])}                changes = []                for key, new_ref in new_refs.items():            old_ref = old_refs.get(key)                        if not old_ref:                changes.append({                    "key": key,                    "change": "added",                    "new_version": new_ref.version                })            elif old_ref.version != new_ref.version:                changes.append({                    "key": key,                    "change": "updated",                    "old_version": old_ref.version,                    "new_version": new_ref.version                })                return changes