Designing Good MCP Servers

"""docs-assistant designed for an LLM caller: few tools, two detail levels, cursors, actionable errors, readable IDs."""
import base64
import json
from typing import Literal

from pydantic import BaseModel, Field

from mcp.server import MCPServer
from mcp.server.mcpserver.exceptions import ToolError
from mcp.types import ToolAnnotations

import store

mcp = MCPServer(
    "docs-assistant", version="1.0.0",
    instructions="Project documentation. Call search_docs first, then get_page for the one or two pages you will quote.",
)
PAGE_SIZE = 2
Component = Literal["algebra", "units", "statistics"]


class Hit(BaseModel):
    id: str = Field(description="Readable page ID such as 'algebra/matrices'; pass it unchanged to get_page")
    title: str


class FullHit(Hit):
    snippet: str
    score: int


class SearchResult(BaseModel):
    hits: list[FullHit | Hit]
    next_cursor: str | None = Field(None, description="Pass as cursor to get the next page; null when done")


def _encode(offset: int) -> str:
    return base64.urlsafe_b64encode(json.dumps({"o": offset}).encode()).decode()


def _decode(cursor: str) -> int:
    try:
        return int(json.loads(base64.urlsafe_b64decode(cursor.encode()))["o"])
    except (ValueError, KeyError, TypeError):
        raise ToolError("Invalid cursor. Use the next_cursor value from the previous result, "
                        "or omit cursor to start from the first page.") from None


@mcp.tool(annotations=ToolAnnotations(readOnlyHint=True, openWorldHint=False))
def search_docs(query: str, component: Component | None = None,
                detail: Literal["concise", "full"] = "concise",
                cursor: str | None = None) -> SearchResult:
    """Find documentation pages. Start with detail='concise' (ID and title only, cheap);
    call get_page for the one or two pages you need to read. If next_cursor is not null,
    pass it as cursor to see more results instead of rewording the query."""
    if not query.strip():
        raise ToolError("query is empty. Pass a key phrase from the question, for example 'LU decomposition'.")
    hits = []
    for c in ([component] if component else store.COMPONENTS):
        hits += store.search(query, c, 50)
    hits.sort(key=lambda h: (-h["score"], h["component"], h["page"]))
    start = _decode(cursor) if cursor else 0
    window = hits[start:start + PAGE_SIZE]
    if not window and start == 0:
        raise ToolError(f"No pages match '{query}'. Try fewer or more general words, or drop the component filter.")
    items = [FullHit(id=f"{h['component']}/{h['page']}", title=h["title"], snippet=h["snippet"], score=h["score"])
             if detail == "full" else Hit(id=f"{h['component']}/{h['page']}", title=h["title"]) for h in window]
    more = start + PAGE_SIZE < len(hits)
    return SearchResult(hits=items, next_cursor=_encode(start + PAGE_SIZE) if more else None)


@mcp.tool(annotations=ToolAnnotations(readOnlyHint=True, openWorldHint=False))
def get_page(id: str) -> str:
    """Return the AsciiDoc source of one page. id comes from search_docs, e.g. 'algebra/matrices'."""
    component, _, page = id.partition("/")
    try:
        return store.read(component, page)
    except KeyError:
        raise ToolError(f"Unknown page '{id}'. IDs look like 'algebra/matrices'; call search_docs to find one.") from None


if __name__ == "__main__":
    mcp.run("streamable-http")

list_components_components_get     create_component_components_post   get_component_components
replace_component_components       delete_component_components        list_pages_components
create_page_components             get_page_pages                     delete_page_pages
update_page_pages                  search_search_get                  list_revisions_pages

"""Estimate the context cost of a server's tool definitions: T = sum(d_i + s_i), at ~4 characters per token."""
import json
import sys

CHARS_PER_TOKEN = 4  # rule of thumb; count with your model provider's tokenizer for a real budget


def budget(tools: list[dict]) -> tuple[int, list[tuple[str, int, int]]]:
    rows = []
    for t in tools:
        d = len(t["name"]) + len(t.get("description") or "")
        s = len(json.dumps(t.get("inputSchema", {}), separators=(",", ":")))
        s += len(json.dumps(t.get("outputSchema", {}), separators=(",", ":")))
        rows.append((t["name"], d // CHARS_PER_TOKEN, s // CHARS_PER_TOKEN))
    return sum(d + s for _, d, s in rows), rows


if __name__ == "__main__":
    tools = json.load(open(sys.argv[1]))["result"]["tools"]
    total, rows = budget(tools)
    for name, d, s in rows:
        print(f"{name:40} d={d:4} s={s:4}")
    print(f"{len(rows)} tools, T = {total} tokens")

search_docs                              d=  63 s= 351
get_page                                 d=  24 s=  59
2 tools, T = 497 tokens

list_components_components_get           d=  11 s=  93
create_component_components_post         d=  12 s=  67
...
12 tools, T = 1047 tokens

import { Client, StreamableHTTPClientTransport } from '@modelcontextprotocol/client';

const client = new Client({ name: 'code-mode-runtime', version: '1.0.0' });
let connected: Promise<void> | undefined;

/** The one place that talks MCP: connection, auth headers and logging live here, not in generated files. */
export async function callMcpTool<T>(tool: string, args: Record<string, unknown>): Promise<T> {
  // Share one connection, but forget a failed attempt so the next call reconnects instead of rethrowing it forever.
  connected ??= client.connect(new StreamableHTTPClientTransport(new URL(process.env.DOCS_MCP_URL!)))
    .catch((err: unknown) => { connected = undefined; throw err; });
  await connected;
  const result = await client.callTool({ name: tool, arguments: args });
  if (result.isError) throw new Error(JSON.stringify(result.content));
  return result.structuredContent as T;
}

export const closeMcp = () => client.close();

import { callMcpTool } from '../../client.js';

export interface SearchDocsInput {
  query: string;
  component?: 'algebra' | 'units' | 'statistics';
  /** 'concise' returns IDs and titles only; 'full' adds snippets and scores. */
  detail?: 'concise' | 'full';
  cursor?: string;
}

export interface Hit { id: string; title: string; snippet?: string; score?: number }
export interface SearchDocsOutput { hits: Hit[]; next_cursor: string | null }

/** Find documentation pages. Follow next_cursor for more results. */
export function searchDocs(input: SearchDocsInput): Promise<SearchDocsOutput> {
  return callMcpTool<SearchDocsOutput>('search_docs', { ...input });
}

// What the model writes: page through results in code, keep only what the answer needs.
import * as docs from './servers/docs/index.js';
import { closeMcp } from './client.js';

const ids: string[] = [];
let cursor: string | undefined;
do {
  const page = await docs.searchDocs({ query: 'matrix solve convert distributions', cursor });
  ids.push(...page.hits.map(h => h.id));
  cursor = page.next_cursor ?? undefined;
} while (cursor);

console.log(`${ids.length} matching pages; first two: ${ids.slice(0, 2).join(', ')}`);   // only this reaches the model
await closeMcp();