# server.py
from mcp.server import Server, stdio_server
from statebase import StateBase

sb = StateBase(api_key="your-key")
server = Server("statebase-mcp")


@server.tool()
async def get_state(session_id: str) -> dict:
    """Fetch the current state for a session."""
    return sb.sessions.get(session_id=session_id).state


@server.tool()
async def save_state(session_id: str, state: dict, reasoning: str = "") -> dict:
    """Persist a checkpoint for a session."""
    return sb.sessions.update_state(
        session_id=session_id, state=state, reasoning=reasoning
    )


@server.tool()
async def recall(session_id: str, query: str) -> list:
    """Semantic recall from agent memory."""
    return sb.memory.search(session_id=session_id, query=query)


async def main():
    async with stdio_server() as (read, write):
        await server.run(read, write)


if __name__ == "__main__":
    import asyncio
    asyncio.run(main())

pip install statebase mcp
python server.py

# An MCP agent that remembers across sessions
from mcp.client import stdio_client
from statebase import StateBase

sb = StateBase(api_key="your-key")
session_id = "mcp-agent-prod"

# 1. Before each run, rehydrate context
context = sb.sessions.get_context(session_id=session_id, query=user_query)

# 2. Send to the MCP-capable agent along with recalled context
messages = [{"role": "system", "content": f"Context: {context}"}] + history

# 3. After the run, checkpoint
sb.sessions.add_turn(
    session_id=session_id,
    input=user_query,
    output=agent_response,
    reasoning="MCP tool interaction logged"
)