{ "cells": [ { "cell_type": "raw", "metadata": { "vscode": { "languageId": "raw" } }, "source": [ "# 🔄 State Management in MCP Tools\n", "\n", "Learn how to effectively manage state in your MCP tools. This notebook covers different state management patterns, persistence strategies, and best practices for maintaining tool state.\n", "\n", "## 🎯 Learning Objectives\n", "\n", "By the end of this notebook, you will:\n", "- Understand different types of tool state\n", "- Implement state persistence\n", "- Handle concurrent state access\n", "- Manage state transitions\n", "- Build stateful MCP tools\n", "\n", "## 📋 Prerequisites\n", "\n", "- Completed notebooks 01-08\n", "- Basic understanding of concurrency\n", "- Python virtual environment with required packages\n", "\n", "## 🔑 Key Concepts\n", "\n", "1. **State Types**\n", " - Memory state\n", " - Persistent state\n", " - Shared state\n", " - Isolated state\n", "\n", "2. **State Management Patterns**\n", " - State machines\n", " - Event sourcing\n", " - Command pattern\n", " - Observer pattern\n", "\n", "3. **Persistence Strategies**\n", " - File-based\n", " - Database\n", " - Cache\n", " - Distributed state\n", "\n", "## 📚 Table of Contents\n", "\n", "1. [Understanding Tool State](#understanding)\n", "2. [State Management Patterns](#patterns)\n", "3. [Persistence Strategies](#persistence)\n", "4. [Concurrent Access](#concurrency)\n", "5. [Best Practices](#practices)\n", "\n", "---\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import json\n", "from typing import Dict, Any, Optional\n", "from dataclasses import dataclass, asdict\n", "from pathlib import Path\n", "from enum import Enum\n", "import threading\n", "import modelcontextprotocol as mcp\n", "from pydantic import BaseModel, Field\n", "\n", "# Define state types\n", "class StateType(Enum):\n", " MEMORY = \"memory\"\n", " FILE = \"file\"\n", " DATABASE = \"database\"\n", "\n", "class StateInput(BaseModel):\n", " key: str = Field(..., description=\"State key to get/set\")\n", " value: Optional[Any] = Field(None, description=\"Value to set (if setting state)\")\n", " \n", "class StateOutput(BaseModel):\n", " key: str\n", " value: Any\n", " state_type: StateType\n", " \n", "class StatefulTool:\n", " def __init__(self, state_type: StateType = StateType.MEMORY):\n", " self.state_type = state_type\n", " self.state: Dict[str, Any] = {}\n", " self._lock = threading.Lock()\n", " \n", " if state_type == StateType.FILE:\n", " self.state_file = Path(\"tool_state.json\")\n", " self._load_state()\n", " \n", " def _load_state(self):\n", " \"\"\"Load state from file if it exists.\"\"\"\n", " if self.state_file.exists():\n", " with open(self.state_file, 'r') as f:\n", " self.state = json.load(f)\n", " \n", " def _save_state(self):\n", " \"\"\"Save state to file if using file persistence.\"\"\"\n", " if self.state_type == StateType.FILE:\n", " with open(self.state_file, 'w') as f:\n", " json.dump(self.state, f)\n", " \n", " def get_state(self, key: str) -> StateOutput:\n", " \"\"\"Get state value for a key.\"\"\"\n", " with self._lock:\n", " value = self.state.get(key)\n", " return StateOutput(\n", " key=key,\n", " value=value,\n", " state_type=self.state_type\n", " )\n", " \n", " def set_state(self, key: str, value: Any) -> StateOutput:\n", " \"\"\"Set state value for a key.\"\"\"\n", " with self._lock:\n", " self.state[key] = value\n", " if self.state_type == StateType.FILE:\n", " self._save_state()\n", " return StateOutput(\n", " key=key,\n", " value=value,\n", " state_type=self.state_type\n", " )\n", " \n", " async def handle_state(self, input_data: StateInput) -> StateOutput:\n", " \"\"\"MCP handler for state operations.\"\"\"\n", " if input_data.value is not None:\n", " return self.set_state(input_data.key, input_data.value)\n", " return self.get_state(input_data.key)\n", "\n", "# Create our MCP tool\n", "state_tool = StatefulTool(state_type=StateType.FILE)\n", "\n", "# Create the MCP server\n", "server = mcp.Server()\n", "server.add_tool(\"state\", state_tool.handle_state, StateInput, StateOutput)\n" ] }, { "cell_type": "raw", "metadata": { "vscode": { "languageId": "raw" } }, "source": [ "# Using the State Management Tool\n", "\n", "Let's test our state management tool with different scenarios:\n", "\n", "1. **Basic State Operations**\n", " - Setting state\n", " - Getting state\n", " - Updating state\n", " - Deleting state\n", "\n", "2. **Persistence**\n", " - Saving to file\n", " - Loading from file\n", " - State recovery\n", "\n", "3. **Concurrent Access**\n", " - Multiple operations\n", " - Thread safety\n", " - Race condition prevention\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Basic state operations\n", "print(\"Setting initial state...\")\n", "result = state_tool.set_state(\"user_preferences\", {\n", " \"theme\": \"dark\",\n", " \"language\": \"en\",\n", " \"notifications\": True\n", "})\n", "print(f\"Set state: {result}\")\n", "\n", "print(\"\\nGetting state...\")\n", "result = state_tool.get_state(\"user_preferences\")\n", "print(f\"Got state: {result}\")\n", "\n", "print(\"\\nUpdating state...\")\n", "preferences = result.value\n", "preferences[\"theme\"] = \"light\"\n", "result = state_tool.set_state(\"user_preferences\", preferences)\n", "print(f\"Updated state: {result}\")\n", "\n", "# Test persistence\n", "print(\"\\nTesting file persistence...\")\n", "new_tool = StatefulTool(state_type=StateType.FILE)\n", "result = new_tool.get_state(\"user_preferences\")\n", "print(f\"Loaded state from file: {result}\")\n", "\n", "# Test concurrent access\n", "import threading\n", "import time\n", "\n", "def concurrent_updates():\n", " for i in range(5):\n", " state_tool.set_state(f\"counter_{threading.current_thread().name}\", i)\n", " time.sleep(0.1)\n", "\n", "print(\"\\nTesting concurrent access...\")\n", "threads = [\n", " threading.Thread(target=concurrent_updates, name=f\"thread_{i}\")\n", " for i in range(3)\n", "]\n", "\n", "for thread in threads:\n", " thread.start()\n", " \n", "for thread in threads:\n", " thread.join()\n", "\n", "print(\"\\nFinal state after concurrent updates:\")\n", "for i in range(3):\n", " result = state_tool.get_state(f\"counter_thread_{i}\")\n", " print(f\"Thread {i} final count: {result.value}\")\n" ] }, { "cell_type": "raw", "metadata": { "vscode": { "languageId": "raw" } }, "source": [ "# Advanced State Management Patterns\n", "\n", "When building production MCP tools, consider these advanced patterns:\n", "\n", "## 1. State Machine Pattern\n", "\n", "Use state machines when your tool has well-defined states and transitions:\n", "- Define valid states\n", "- Control state transitions\n", "- Validate state changes\n", "- Handle edge cases\n", "\n", "## 2. Event Sourcing\n", "\n", "Track all state changes as a series of events:\n", "- Record every state modification\n", "- Enable state replay\n", "- Support audit trails\n", "- Allow time travel debugging\n", "\n", "## 3. Command Pattern\n", "\n", "Encapsulate state changes as commands:\n", "- Separate state modification logic\n", "- Enable undo/redo\n", "- Support command queuing\n", "- Implement command validation\n", "\n", "## 4. Observer Pattern\n", "\n", "Notify components of state changes:\n", "- Subscribe to state updates\n", "- React to state changes\n", "- Maintain consistency\n", "- Decouple components\n", "\n", "## Best Practices\n", "\n", "1. **State Validation**\n", " - Validate state before saving\n", " - Check state integrity\n", " - Handle invalid states\n", " - Maintain state invariants\n", "\n", "2. **Error Handling**\n", " - Handle persistence failures\n", " - Recover from corruption\n", " - Maintain consistency\n", " - Log state changes\n", "\n", "3. **Performance**\n", " - Cache frequent access\n", " - Batch state updates\n", " - Optimize persistence\n", " - Monitor state size\n", "\n", "4. **Security**\n", " - Encrypt sensitive state\n", " - Control access\n", " - Audit state changes\n", " - Sanitize inputs\n", "\n", "## Exercise\n", "\n", "Implement these improvements to our state tool:\n", "\n", "1. Add state validation\n", "2. Implement event sourcing\n", "3. Add command pattern support\n", "4. Create state observers\n" ] }, { "cell_type": "raw", "metadata": { "vscode": { "languageId": "raw" } }, "source": [ "# 🔄 State Management in MCP\n", "\n", "Master stateful MCP applications! Learn to manage persistent state, user sessions, workflow orchestration, and build MCP servers that remember context across interactions.\n", "\n", "## 🎯 Learning Objectives\n", "\n", "By the end of this notebook, you will:\n", "- Implement persistent state storage patterns\n", "- Manage user sessions and context\n", "- Build stateful workflows and multi-step processes\n", "- Handle concurrent state updates safely\n", "- Design scalable state architectures\n", "\n", "## 🛠️ State Management Patterns\n", "\n", "- **📝 Session State**: User-specific temporary data\n", "- **💾 Persistent State**: Long-term data storage\n", "- **🔄 Workflow State**: Multi-step process management\n", "- **🔐 Shared State**: Thread-safe concurrent access\n", "\n", "## 📚 Table of Contents\n", "\n", "1. [State Storage Backends](#storage)\n", "2. [Session Management](#sessions)\n", "3. [Workflow Orchestration](#workflows)\n", "4. [Concurrency & Thread Safety](#concurrency)\n", "5. [State Synchronization](#synchronization)\n", "\n", "---\n" ] } ], "metadata": { "language_info": { "name": "python" } }, "nbformat": 4, "nbformat_minor": 2 }