Lesson 09 Authentication & Secrets
Theme:

Solutions — Lesson 09

Reference implementations with the most common mistakes called out

Don't read these until you've attempted the tasks. Each solution is one of several reasonable answers — the goal is to highlight the discipline (per-user buckets, redact at the boundary, soft-confirm destructive actions), not to enforce a particular spelling.

Solution to Task 1 — Per-user write

The two changes are: read the user with get_current_user(ctx), then use a setdefault'd per-user bucket for both the existence check and the write. 

@mcp.tool()
def write(args: WriteArgs, ctx: Context) -> str:
    """Create or overwrite a note for the current user."""
    user = get_current_user(ctx)
    bucket = USER_NOTES.setdefault(user, {})
    if args.title in bucket and not args.overwrite:
        raise McpError(ErrorData(code=INVALID_PARAMS,
            message=f"note exists: {args.title}; pass overwrite=true to replace"))
    bucket[args.title] = args.body
    return f"saved: {args.title}"

Common mistake

Forgetting that the module-level NOTES alias still points at USER_NOTES["local"] — the bucket created at import time. If you keep writing to NOTES you're always writing to Alice's notes when PKB_USER=alice only by accident (because she happens to be "local"); Bob's writes will silently land in Alice's bucket. The fix is to always resolve the bucket from USER_NOTES[user] on every call.

Solution to Task 2 — Secret hygiene audit

The redaction lives in a tiny helper that the error formatter calls. Reading the key once at import time is fine for a dev stub; in production you'd resolve it through a secrets manager on each call so rotation works without a restart. 

THIRD_PARTY_KEY = os.environ.get("PKB_TRANSLATE_KEY", "")


def _redact(text: str) -> str:
    return text.replace(THIRD_PARTY_KEY, "[REDACTED]") if THIRD_PARTY_KEY else text


@mcp.tool()
def translate(id: str, lang: str) -> str:
    """Translate a note via a hypothetical third-party API."""
    try:
        resp = _call_translate_api(NOTES[id], lang)  # may raise
        return resp.text
    except Exception as e:
        raise McpError(ErrorData(
            code=INTERNAL_ERROR,
            message=f"translate failed: {_redact(str(e))}",
        ))

And the test:

def test_no_key_leak(monkeypatch):
    monkeypatch.setenv("PKB_TRANSLATE_KEY", "secret-token-xyz")
    monkeypatch.setattr("pkb.tools._call_translate_api",
                        lambda *a, **k: (_ for _ in ()).throw(
                            RuntimeError("GET https://api.example/?key=secret-token-xyz failed")))
    with pytest.raises(McpError) as exc:
        translate(id="hello", lang="es")
    assert "secret-token-xyz" not in str(exc.value)
    assert "[REDACTED]" in str(exc.value)

Common mistake

Stripping the key only from the top-level message but not from the data field of ErrorData. If you ever attach the original exception (or any string derived from it) to data, that path needs redaction too. The defensive move: redact at the boundary — one helper, one call site, no exceptions. Bonus: if your test asserts on str(exc.value) it catches both surfaces in one line.

Solution to Task 3 — Lock down export

A soft return is the right shape here. The model gets a clear, readable instruction and the transcript records exactly what happened; an error frame would suggest something went wrong. 

@mcp.tool()
def export(ctx: Context, confirm_export: bool = False) -> dict[str, str] | str:
    """Export all of the current user's notes. Requires confirm_export=true."""
    if not confirm_export:
        return "set confirm_export=true to proceed (destructive disclosure)"
    user = get_current_user(ctx)
    return export_notes(user)

Common mistake

Raising McpError instead of returning a plain string. For a user-facing “did you really mean it?” check, a soft return reads better than a hard error. Save errors for protocol violations (missing auth, malformed input, downstream API failure) and use return values for everything the model is expected to act on.

Where this goes next

Module 4 (Plan D) introduces structured logging — including how to not log secrets — and a real testing harness for sampling tools. The redaction helper from Task 2 graduates into a proper log filter that scrubs configured secret names from every emitted log line, and the per-user bucket from Task 1 becomes the unit of isolation in the testing harness.