Lab 5: Build a Chat Client With MCP¶

"""Lab 5: a terminal chat client that uses tools from any MCP server.

    python chat.py                  # servers from mcp_servers.json
Commands: /tools  /tokens  /clear  /quit
"""

import asyncio
import json
import os
from contextlib import AsyncExitStack

import anthropic
from mcp import ClientSession
from mcp.client.stdio import StdioServerParameters, stdio_client

MODEL = os.environ.get("AGENT_MODEL", "claude-opus-5-5")
MAX_STEPS = 10
SYSTEM = "You are a helpful on-call assistant. Use the tools when a question needs evidence. Be brief."

client = anthropic.AsyncAnthropic()


async def connect(stack, config_path="mcp_servers.json"):
    """Start every configured server and return {tool_name: session} and Claude tool definitions."""
    config = json.loads(open(config_path).read())["mcpServers"]
    routes, tools = {}, []
    for name, server in config.items():
        params = StdioServerParameters(command=server["command"], args=server.get("args", []),
                                       env={**os.environ, **server.get("env", {})})
        read, write = await stack.enter_async_context(stdio_client(params))
        session = await stack.enter_async_context(ClientSession(read, write))
        await session.initialize()
        for tool in (await session.list_tools()).tools:
            if tool.name in routes:
                print(f"skipping {tool.name} from {name}: another server already defines it")
                continue
            routes[tool.name] = session
            tools.append({"name": tool.name, "description": tool.description or "",
                          "input_schema": tool.input_schema})
        print(f"connected to {name}: {', '.join(t['name'] for t in tools if routes[t['name']] is session)}")
    return routes, tools


async def call_mcp_tool(session, name, args):
    result = await session.call_tool(name, args)
    text = "\n".join(part.text for part in result.content if part.type == "text")
    return text or "(no text output)", bool(result.is_error)


async def answer(messages, routes, tools):
    """Run the tool loop for one user message. Appends every turn to `messages`."""
    for _ in range(MAX_STEPS):
        response = await client.messages.create(
            model=MODEL, max_tokens=16000, system=SYSTEM, tools=tools, messages=messages,
            cache_control={"type": "ephemeral"},
        )
        messages.append({"role": "assistant", "content": response.content})
        if response.stop_reason != "tool_use":
            return "".join(b.text for b in response.content if b.type == "text")
        results = []
        for block in response.content:
            if block.type == "tool_use":
                print(f"  [{block.name} {json.dumps(block.input)}]")
                content, is_error = await call_mcp_tool(routes[block.name], block.name, block.input)
                results.append({"type": "tool_result", "tool_use_id": block.id,
                                "content": content, "is_error": is_error})
        messages.append({"role": "user", "content": results})
    return f"(stopped after {MAX_STEPS} tool steps)"


async def main():
    async with AsyncExitStack() as stack:
        routes, tools = await connect(stack)
        messages = []  # the conversation: this is the chat's whole short-term memory
        while True:
            try:
                text = (await asyncio.to_thread(input, "\nyou> ")).strip()
            except EOFError:
                break
            if text in ("/quit", "/exit"):
                break
            if text == "/tools":
                print("\n".join(f"  {t['name']}: {t['description']}" for t in tools))
                continue
            if text == "/clear":
                messages = []
                print("  conversation cleared")
                continue
            if text == "/tokens":
                # A request must end with a user turn, so count with a placeholder next message.
                count = await client.messages.count_tokens(
                    model=MODEL, system=SYSTEM, tools=tools,
                    messages=messages + [{"role": "user", "content": "(next message)"}])
                print(f"  {count.input_tokens:,} tokens of context go out with your next message")
                continue
            if not text:
                continue
            messages.append({"role": "user", "content": text})
            print(f"\nclaude> {await answer(messages, routes, tools)}")


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