Lab 09 — Add an OAuth-gated export tool

Lesson 09 · MCP Servers · ~30 min hands-on

By the end of this lab the PKB server exposes a new export tool that requires authentication. You'll see it succeed when PKB_USER is set in the environment and fail with a structured error when it isn't. Along the way you'll learn the convention this course uses for handling secrets — .env files, gitignored, never echoed in error messages — and how to wire python-dotenv into the server entry point for production.

Prereqs

• Lesson 08 lab done — PKB exposes 5 tools (echo, search, summarize, fetch, write, summarize_and_save) and 2 resources (pkb://about, pkb://notes/{title}).
• Editable install still in place from Lesson 08: pip install -e . from mcp-servers/sample-server/python/.
• python-dotenv available: pip install python-dotenv.

Step 1 — Create pkb/auth.py

Create a new file mcp-servers/sample-server/python/pkb/auth.py. This module is the integration point between PKB and a real OAuth provider. For lessons 09–10 the body is a stub that reads from an env var; in a remote deployment (Module 4 / Lessons 13–14) you'd swap the body to extract the sub claim from a validated OAuth access token off ctx.request_context. The public surface — get_current_user(ctx) -> str raising McpError(INVALID_PARAMS) on missing auth — stays identical.

"""Stub auth integration for PKB.

This module is the integration point between PKB and a real OAuth provider.
The function `get_current_user(ctx)` returns the authenticated user id.

For lessons 09-10 we read the user from `PKB_USER` env var (single-process
dev). In a remote deployment (Lessons 13-14), swap the body to extract the
`sub` claim from a validated OAuth access token off `ctx.request_context`.
"""

from __future__ import annotations

import os

from mcp.shared.exceptions import McpError
from mcp.types import ErrorData, INVALID_PARAMS


def get_current_user(ctx) -> str:
    """Return the authenticated user id, or raise McpError if missing.

    Stub implementation. Real impl: extract claim from validated OAuth token.
    """
    user = os.environ.get("PKB_USER", "").strip()
    if not user:
        raise McpError(ErrorData(
            code=INVALID_PARAMS,
            message="not authenticated (PKB_USER env var not set)",
        ))
    return user

What each piece is doing:

• Module docstring. Calls out explicitly that this is the seam for swapping in real OAuth later. A future reader (you, in three months) won't have to guess where the integration boundary lives.
• get_current_user(ctx). The only public function. Every tool that needs auth calls this — never reads os.environ directly. That centralises the stub-vs-real decision in one place.
• McpError(ErrorData(code=INVALID_PARAMS, ...)). We raise a structured MCP error rather than a generic Python exception. The protocol code (-32602) and the message both land in the client's error UI with no extra plumbing. In a real implementation you'd map an invalid/expired token to the same shape.

Step 2 — Extend pkb/tools.py

Open pkb/tools.py and refactor the top of the file. The single-tenant NOTES dict from Lessons 04–08 becomes a view into a multi-user store, plus a small helper for the export tool:

"""Tool implementations for PKB. Imported by server.py."""

from pydantic import BaseModel, Field
from mcp.shared.exceptions import McpError
from mcp.types import ErrorData, INVALID_PARAMS

# Per-user note store. NOTES (legacy name) views into the default user's bucket
# so prior lesson code keeps working in the lab.
USER_NOTES: dict[str, dict[str, str]] = {
    "local": {
        "hello": "Hello note body",
        "todo": "Pick up groceries",
        "ideas": "Build an MCP server tutorial",
    }
}

# Legacy single-tenant alias used by L04-L08 examples. Points at the "local" bucket.
NOTES = USER_NOTES["local"]


def export_notes(user: str) -> dict[str, str]:
    """Return a copy of all notes for USER. Used by the export tool."""
    return dict(USER_NOTES.get(user, {}))

What's going on here:

• USER_NOTES. A two-level dict keyed by user id, then by note title. This is the minimal shape for per-tenant state — small enough to fit in memory for the lab, large enough to expose the auth boundary clearly.
• NOTES = USER_NOTES["local"]. A live view into the "local" bucket. Because Python dicts are mutable references, every tool from Lessons 04–08 that reads or writes NOTES keeps working unchanged — it just now reads and writes the "local" user's bucket. This is the bridge that lets us introduce multi-tenancy without breaking the existing lab surface.
• export_notes(user). Returns a copy of the user's notes (not the live dict) so the tool can't be used to mutate state through its return value. Falls back to an empty dict for unknown users — the auth gate is what guarantees we never get here with a bogus user.

Step 3 — Wire export into server.py

Open pkb/server.py. Add the new imports near the existing tool imports, then register the export tool after the others:

from pkb.tools import export_notes
from pkb.auth import get_current_user


@mcp.tool()
def export(ctx: Context) -> dict[str, str]:
    """Export all of the current user's notes. Requires authentication."""
    user = get_current_user(ctx)
    return export_notes(user)

Two things to notice. First, the tool body is a two-liner: ask auth who the caller is, then ask tools for that user's notes. The tool itself contains no auth logic — it just composes the two. Second, get_current_user(ctx) raises before export_notes is ever called, so the USER_NOTES.get(user, {}) fallback never sees an unauthenticated request.

Step 4 — Add a .env file

Create mcp-servers/sample-server/python/.env with the dev credentials:

PKB_USER=local
# Future: PKB_THIRD_PARTY_KEY=...

The repo's root .gitignore already excludes .env files. Confirm before you go further — committing a populated .env is the single most common way developers leak secrets to public repos:

git check-ignore mcp-servers/sample-server/python/.env

Expected output: the path you passed in, echoed back. If you get an empty line, stop and add .env to .gitignore before continuing — never commit a file you intend to put secrets in.

Step 5 — Load .env at startup (optional)

For this lab we'll pass PKB_USER inline on the mcp dev command, which is the cleanest way to see the auth path under both states (set / unset). For production deployments, wire python-dotenv into the top of pkb/server.py so a real run picks up the .env automatically:

from dotenv import load_dotenv
load_dotenv()

Note: mcp dev does not load .env for you — that's a property of the mcp[cli] harness, not Python in general. If you add the two lines above, mcp dev will still work, and a plain python -m pkb.server will pick the env up automatically.

Step 6 — Reinstall and run with auth

The new auth.py module is a new file, so a fresh editable install is needed to pick it up. We also install python-dotenv for Step 5's optional wiring:

cd mcp-servers/sample-server/python
pip install -e . python-dotenv
PKB_USER=local mcp dev pkb/server.py

Inspector opens. Click the Tools tab — you should now see seven entries: the six from Lesson 08 (echo, search, summarize, fetch, write, summarize_and_save) plus the new export. If export is missing, scroll the terminal back — FastMCP logs import failures before Inspector even loads, so a broken import in auth.py or a typo in the tool registration usually shouts there first.

Step 7 — Call export (authenticated)

In Inspector, click Tools → export and submit with no arguments. Expected return value:

{
  "hello": "Hello note body",
  "todo": "Pick up groceries",
  "ideas": "Build an MCP server tutorial"
}

That's the local user's bucket from USER_NOTES — exactly the same three seed notes you've been reading and writing throughout the course. The auth gate let the call through because PKB_USER=local matched a real key in USER_NOTES.

Step 8 — Call export (unauthenticated)

Stop the server (Ctrl-C in the terminal). Restart it without the env var:

mcp dev pkb/server.py

Back in Inspector (re-open it if it disconnected), call export again with no arguments. Expected: a structured MCP error, not a Python traceback:

code:    -32602  (INVALID_PARAMS)
message: not authenticated (PKB_USER env var not set)

The auth gate held. The error code (-32602) is a real protocol error — the client UI treats it as a tool failure and surfaces it to the user, not as a server crash. The message names the variable (PKB_USER) but never echoes its value. That's a small but deliberate choice: an error path that leaks the secret it was checking would be worse than no check at all.

Step 9 — Inspect for credential leaks

Before you call the lab done, walk through this checklist against your own code. The PKB server passes all three; the point of the exercise is to internalise the pattern so your next server passes too.

What you should leave with

• PKB exposes 6 tools now (Lesson 08's five plus the new export), and one of them — export — is gated by authentication.
• Auth lives in one place: pkb/auth.py. Every tool that needs the caller's identity calls get_current_user(ctx); nothing else reads os.environ for credentials directly. That makes the stub-to-real swap in Module 4 a single-file change.
• The unauthenticated path returns a structured MCP error (-32602 INVALID_PARAMS), not a stack trace. Clients can render it; LLMs can recover from it; humans can read it.
• .env handling is now muscle memory: gitignored, loadable via python-dotenv for production, never echoed in error messages.

Troubleshooting