Lab 10 — Instrument PKB and run its first test suite

Lesson 10 · MCP Servers · ~40 min hands-on

By the end of this lab, every PKB tool call emits a structured log line (phase=start, then phase=ok dur_ms=… or phase=err type=…); no tool arguments and no LLM sampling output appear in the logs; and pytest pkb/tests/ passes locally against the in-memory MCP client/server harness. You are giving PKB the two cheap habits that let Modules 5–6 (and your future self) change it without breaking it: logs you can read and tests that fail when behaviour drifts.

Prereqs

• Plan C / Lesson 09 lab done. PKB exposes 7 tools (echo, search, summarize, fetch, write, summarize_and_save, export), 2 resources (pkb://about, pkb://notes/{title}), 1 prompt (welcome), with auth.py in place.
• Editable install still active from Lesson 09: pip install -e . from mcp-servers/sample-server/python/.
• Install dev deps for the test suite:

  pip install pytest pytest-asyncio

  (If you added the [project.optional-dependencies] dev = […] block to pyproject.toml, you can instead run pip install -e ".[dev]".)

Step 1 — Create pkb/observability.py

Create the file mcp-servers/sample-server/python/pkb/observability.py. This module is PKB's single source of logging configuration, the per-tool log decorator, the trace-id generator, and the redaction helper. Every other module pulls from it — no per-tool logging boilerplate elsewhere.

"""Observability primitives for PKB.

Provides:
- _configure_logging(): single-call logging setup used by server.py at startup.
- log_tool_call(name): decorator that wraps any tool function with start/ok/err
  log lines and a duration in milliseconds. Works with both sync and async tools.
- redact(value): mask helper for fields that might contain user data or secrets.
- new_trace_id(): short random id for per-tool tracing.
"""

from __future__ import annotations

import asyncio
import functools
import logging
import secrets
import time
from typing import Any, Callable


def _configure_logging(level: int = logging.INFO) -> None:
    """Idempotent logging setup. Safe to call multiple times."""
    if logging.getLogger().handlers:
        return
    logging.basicConfig(
        level=level,
        format="%(asctime)s %(levelname)s %(name)s %(message)s",
    )


def new_trace_id() -> str:
    """8-char hex trace id — short enough to scan in logs, unique enough for a session."""
    return secrets.token_hex(4)


def redact(value: Any, max_len: int = 200) -> str:
    """Return a safe-to-log representation of `value`."""
    if isinstance(value, str):
        if len(value) > max_len:
            return f"<str len={len(value)} preview={value[:32]!r}...>"
        return value
    if isinstance(value, (bytes, bytearray)):
        return f"<bytes len={len(value)}>"
    s = repr(value)
    if len(s) > max_len:
        return f"<{type(value).__name__} len={len(s)}>"
    return s


def log_tool_call(tool_name: str) -> Callable:
    """Decorator: log start / ok / err for the wrapped tool function."""
    def deco(fn: Callable) -> Callable:
        log = logging.getLogger(f"pkb.tool.{tool_name}")

        if asyncio.iscoroutinefunction(fn):
            @functools.wraps(fn)
            async def awrapper(*args, **kwargs):
                trace = new_trace_id()
                log.info("tool=%s trace=%s phase=start", tool_name, trace)
                t0 = time.monotonic()
                try:
                    result = await fn(*args, **kwargs)
                    log.info(
                        "tool=%s trace=%s phase=ok dur_ms=%.1f result_type=%s",
                        tool_name, trace, (time.monotonic() - t0) * 1000,
                        type(result).__name__,
                    )
                    return result
                except Exception as e:
                    log.warning(
                        "tool=%s trace=%s phase=err dur_ms=%.1f type=%s",
                        tool_name, trace, (time.monotonic() - t0) * 1000,
                        type(e).__name__,
                    )
                    raise
            return awrapper

        @functools.wraps(fn)
        def swrapper(*args, **kwargs):
            trace = new_trace_id()
            log.info("tool=%s trace=%s phase=start", tool_name, trace)
            t0 = time.monotonic()
            try:
                result = fn(*args, **kwargs)
                log.info(
                    "tool=%s trace=%s phase=ok dur_ms=%.1f result_type=%s",
                    tool_name, trace, (time.monotonic() - t0) * 1000,
                    type(result).__name__,
                )
                return result
            except Exception as e:
                log.warning(
                    "tool=%s trace=%s phase=err dur_ms=%.1f type=%s",
                    tool_name, trace, (time.monotonic() - t0) * 1000,
                    type(e).__name__,
                )
                raise
        return swrapper

    return deco

What each piece is doing:

• _configure_logging(). Idempotent — the early if logging.getLogger().handlers: return check makes it safe to call multiple times from multiple import points without doubling up handlers. Called once at module import from server.py; emits to stdout with the standard four-field format (time | level | logger | message).
• new_trace_id(). Eight hex characters from secrets.token_hex(4) — long enough to be unique within a session you can hold in your head, short enough to scan visually in a wall of log lines. The composed tool in the homework will use this to thread parent/child trace ids.
• redact(value). The redaction helper. Strings longer than max_len get truncated to <str len=N preview=…>. Bytes get summarised by length. Everything else falls back to repr and is length-checked. Use this whenever you log a field that might contain user data.
• log_tool_call(name). The decorator. Emits at most two lines per call: one phase=start on entry, one phase=ok or phase=err on exit with the duration in milliseconds. Detects async vs sync via asyncio.iscoroutinefunction and chooses the right wrapper. Critically, it logs only metadata — tool name, trace id, phase, duration, result type, exception type — never the arguments themselves and never the return value.

Step 2 — Wire _configure_logging into pkb/server.py

Open pkb/server.py. Add two lines near the top, before any @mcp.tool() decorator runs:

from pkb.observability import _configure_logging, log_tool_call

_configure_logging()

Calling _configure_logging() at the module body means it runs once, at import time, before any tool registration and before any client connects. By the time FastMCP starts wiring tools, the root logger already has a handler and a format. If you defer this to a main() or to the first request, your startup-time tool registrations will log to a half-configured root logger (or, worse, get silently dropped).

Step 3 — Stack @log_tool_call("name") on every tool

For each @mcp.tool() registration in server.py, stack @log_tool_call("<name>") between @mcp.tool() and the function definition:

@mcp.tool()
@log_tool_call("echo")
def echo(text: str) -> str:
    return text

Order matters. @mcp.tool() must be the outermost decorator — it inspects the underlying signature to generate FastMCP's tool schema. @log_tool_call wraps the function before FastMCP sees it, but because the wrapper uses functools.wraps(fn), the wrapped function carries the original __name__, __doc__, and signature. FastMCP's schema generator sees the original signature; the tool's runtime is the wrapped version that emits the log lines.

Apply the decoration to all 7 tools:

• @log_tool_call("echo") on echo
• @log_tool_call("search") on search
• @log_tool_call("summarize") on summarize
• @log_tool_call("fetch") on fetch
• @log_tool_call("write") on write
• @log_tool_call("summarize_and_save") on summarize_and_save
• @log_tool_call("export") on export

Double-check that the string argument to log_tool_call matches the function name. The logger name is built from it (pkb.tool.<name>), and a mismatch will make grepping by tool harder to reason about later.

Step 4 — Run the server and watch the logs

Boot the server through mcp dev so you can poke it from Inspector while watching the log stream:

cd mcp-servers/sample-server/python
PKB_USER=local mcp dev pkb/server.py

Inspector opens in your browser. In the Tools tab, call echo with arguments {"text": "smoke"}. Switch back to your terminal — you should see two new log lines:

2026-05-25 12:00:01,123 INFO pkb.tool.echo tool=echo trace=a1b2c3d4 phase=start
2026-05-25 12:00:01,124 INFO pkb.tool.echo tool=echo trace=a1b2c3d4 phase=ok dur_ms=0.3 result_type=str

Now exercise the error path. Call fetch with {"id": "no-such-note"}. Expected output in the log stream:

2026-05-25 12:00:05,567 INFO pkb.tool.fetch tool=fetch trace=e5f6a7b8 phase=start
2026-05-25 12:00:05,568 WARNING pkb.tool.fetch tool=fetch trace=e5f6a7b8 phase=err dur_ms=0.5 type=McpError

Two things to notice. First, the trace ids (a1b2c3d4, e5f6a7b8) tie the entry line to its corresponding exit line. Second — and this is the entire point of the redaction policy — nowhere does the log contain the literal text "smoke" or the literal id "no-such-note". Only metadata: tool name, trace id, phase, duration, result type, exception type.

Step 5 — Create the pkb/tests/ package

Create an empty file pkb/tests/__init__.py to mark the directory as a Python package. Then create the two test modules below.

pkb/tests/test_tools.py

"""Unit tests for PKB tools via the in-memory MCP client/server harness."""

from __future__ import annotations

import pytest

from mcp.shared.memory import create_connected_server_and_client_session

from pkb import tools as pkb_tools
from pkb.server import mcp


@pytest.fixture
def fresh_notes(monkeypatch):
    """Each test gets its own copy of the seed notes so writes don't leak across tests."""
    seed = {"hello": "Hello note body", "todo": "Pick up groceries",
            "ideas": "Build an MCP server tutorial"}
    monkeypatch.setattr(pkb_tools, "NOTES", dict(seed))
    return seed


@pytest.mark.asyncio
async def test_echo_round_trip():
    async with create_connected_server_and_client_session(mcp._mcp_server) as client:
        result = await client.call_tool("echo", {"text": "hello world"})
        assert result.content[0].text == "hello world"
        assert result.isError is False


@pytest.mark.asyncio
async def test_search_returns_known_titles(fresh_notes):
    async with create_connected_server_and_client_session(mcp._mcp_server) as client:
        result = await client.call_tool("search", {"args": {"query": "hello", "limit": 5}})
        assert "hello" in result.content[0].text


@pytest.mark.asyncio
async def test_fetch_known_id(fresh_notes):
    async with create_connected_server_and_client_session(mcp._mcp_server) as client:
        result = await client.call_tool("fetch", {"id": "hello"})
        assert "Hello note body" in result.content[0].text
        assert result.isError is False


@pytest.mark.asyncio
async def test_fetch_unknown_id_is_error(fresh_notes):
    async with create_connected_server_and_client_session(mcp._mcp_server) as client:
        result = await client.call_tool("fetch", {"id": "no-such-note"})
        assert result.isError is True
        assert "unknown note id" in result.content[0].text


@pytest.mark.asyncio
async def test_write_then_fetch(fresh_notes):
    async with create_connected_server_and_client_session(mcp._mcp_server) as client:
        written = await client.call_tool(
            "write",
            {"args": {"title": "lab10", "body": "from L10 unit test"}},
        )
        assert "saved: lab10" in written.content[0].text

        fetched = await client.call_tool("fetch", {"id": "lab10"})
        assert "from L10 unit test" in fetched.content[0].text


@pytest.mark.asyncio
async def test_summarize_uses_sampling_helper(monkeypatch, fresh_notes):
    """Patch the sampling-dependent helper so the test doesn't need a real LLM."""
    async def fake(note_id, ctx):
        return f"SUMMARY OF {note_id}"
    monkeypatch.setattr(pkb_tools, "summarize_via_sampling", fake)

    async with create_connected_server_and_client_session(mcp._mcp_server) as client:
        result = await client.call_tool("summarize", {"id": "hello"})
        assert result.content[0].text == "SUMMARY OF hello"

pkb/tests/test_resources.py

"""Unit tests for PKB resources via the in-memory client/server harness."""

from __future__ import annotations

import pytest

from mcp.shared.memory import create_connected_server_and_client_session

from pkb import tools as pkb_tools
from pkb.server import mcp


@pytest.fixture
def fresh_notes(monkeypatch):
    seed = {"hello": "Hello note body", "todo": "Pick up groceries",
            "ideas": "Build an MCP server tutorial"}
    monkeypatch.setattr(pkb_tools, "NOTES", dict(seed))
    return seed


@pytest.mark.asyncio
async def test_about_resource_returns_text():
    async with create_connected_server_and_client_session(mcp._mcp_server) as client:
        result = await client.read_resource("pkb://about")
        assert any("pkb" in c.text.lower() for c in result.contents)


@pytest.mark.asyncio
async def test_notes_resource_known(fresh_notes):
    async with create_connected_server_and_client_session(mcp._mcp_server) as client:
        result = await client.read_resource("pkb://notes/hello")
        assert any("Hello note body" in c.text for c in result.contents)


@pytest.mark.asyncio
async def test_notes_resource_unknown(fresh_notes):
    async with create_connected_server_and_client_session(mcp._mcp_server) as client:
        with pytest.raises(Exception):
            await client.read_resource("pkb://notes/does-not-exist")

The patterns worth internalising:

• fresh_notes fixture. Uses monkeypatch.setattr(pkb_tools, "NOTES", dict(seed)) so each test gets its own copy of the seed dict. Writes from one test cannot leak to the next — the fixture's teardown restores the original attribute. The fixture is duplicated in both test modules (rather than promoted to conftest.py) to keep the lab walk-through linear; promoting it is a fair homework variation.
• create_connected_server_and_client_session(mcp._mcp_server). The in-memory harness. It wires a client to a server through a memory channel — no subprocess, no sockets, no transport overhead. Pure Python async. The leading underscore on _mcp_server reaches the low-level server object beneath FastMCP; this access is stable enough in the pinned SDK version but is the most likely thing to drift on an SDK upgrade (see Troubleshooting).
• Sampling is mocked via monkeypatch on summarize_via_sampling, not by passing a sampling callback to the client. Both approaches work; the monkeypatch approach is simpler, keeps the test free of SDK-callback-signature concerns, and produces the same coverage for the tool's own logic.
• test_fetch_unknown_id_is_error checks result.isError is True rather than catching an exception. Tool errors come back from the server as a CallToolResult with isError=True and the error text in the first content block. They do not raise as Python exceptions on the client side — that's a common mistake when porting test patterns from other RPC frameworks.
• Resource errors do raise. The test_notes_resource_unknown test uses pytest.raises(Exception) because the SDK maps an unknown resource URI to a protocol-level error that surfaces as an exception on the client side. Tool errors and resource errors take different paths through the protocol.

Step 6 — Run the tests

cd mcp-servers/sample-server/python
pytest pkb/tests/ -q

Expected output: a row of dots followed by something like 9 passed in 0.42s. The whole suite should finish in well under a second — the in-memory harness has no transport overhead, and the mocked sampling call returns instantly.

If test_summarize_uses_sampling_helper fails because of an SDK signature mismatch on summarize_via_sampling, mark it skipped with a clear reason and continue — see the Troubleshooting box below for the exact fix path.

Step 7 — Read the audit checklist

Verify your own code passes the audit. From the sample-server tree:

cd mcp-servers/sample-server/python
grep -nE "log\.(info|warning|debug|error)\(" pkb/*.py pkb/tests/*.py

Read every match. Every log.…() call should pass only metadata — names, lengths, types, ids, durations — never values. The PKB code shipped in Task 4 already passes this audit; the point of running it manually is to internalise the pattern so your next server passes it too.

What you should leave with

• PKB self-reports every tool call in a redacted, grep-friendly format. Two log lines per successful call, one per failed call; the trace id ties entry to exit. You can answer "what did the server do in the last minute" by piping the log through grep tool=.
• The test suite is a tripwire, not coverage theatre. It asserts the visible behaviour of every tool and resource. If anything drifts — an error message changes, a tool's return type shifts, a resource URI pattern moves — the suite catches it before Inspector ever does.
• pytest pkb/tests/ is a CI-ready gate: no subprocess, no live LLM, no network calls. It runs anywhere Python 3.10+ runs in under a second.
• You have a working mental model for two patterns the production half of the course (Module 5) will build on: structured logging as the seam for OpenTelemetry, and the in-memory client/server harness as the seam for contract tests.

Troubleshooting