Solutions — Lesson 10

Reference answers for Lesson 10's three homework tasks · trace propagation, parametrised error tests, redaction policy

These are reference solutions. Yours may diverge stylistically — what matters is that each self-check passes and the policy document is honest about what your code does. Read the "common mistakes" note after each solution; they capture the failure modes that tripped up reviewers of earlier drafts of this lesson.

Solution to Task 1 — Trace propagation

The trick is to separate roles: the composed tool owns the parent id; the inner call generates its own trace id and additionally records the parent. The log_tool_call decorator keeps doing its job unmodified; we just add one manual log line per function that carries the parent id field.

# In pkb/tools.py — add an optional parent_trace param to the helper:
async def summarize_via_sampling(
    note_id: str,
    ctx: Context,
    parent_trace: str | None = None,
) -> str:
    log = logging.getLogger("pkb.tool.summarize_via_sampling")
    log.info(
        "tool=summarize_via_sampling parent_trace=%s phase=start",
        parent_trace or "-",
    )
    # ... existing body that calls ctx.session.create_message(...) ...
    return summary_text


# In pkb/tools.py — thread the parent trace through the composed tool:
async def summarize_and_save(note_id: str, ctx: Context) -> str:
    from pkb.observability import new_trace_id
    parent = new_trace_id()
    log = logging.getLogger("pkb.tool.summarize_and_save")
    log.info("tool=summarize_and_save trace=%s phase=start", parent)
    summary = await summarize_via_sampling(note_id, ctx, parent_trace=parent)
    new_title = f"{note_id}-summary"
    NOTES[new_title] = summary
    log.info(
        "tool=summarize_and_save trace=%s phase=ok new_title_len=%d",
        parent, len(new_title),
    )
    return f"saved: {new_title}"

The decorator still wraps both functions. So in practice a single call to summarize_and_save emits four decorator lines plus the two manual lines above — six total — with parent_trace= appearing on exactly one of them. That is acceptable: the manual line is the one you actually grep on when correlating an outer call with its inner one.

Solution to Task 2 — Parametrised error cases

# pkb/tests/test_errors.py
"""Parametrised: every PKB tool that takes an id should error cleanly on unknown ids."""

from __future__ import annotations

import pytest

from mcp.shared.memory import create_connected_server_and_client_session
from pkb.server import mcp


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

Note the decorator order: @pytest.mark.parametrize on top, @pytest.mark.asyncio closest to the function. Swapping them works in current pytest-asyncio versions but the documented order is this one and the failure mode (silently skipped tests) is too costly to risk.

About sampling mocking: all three tools perform their "unknown id" check before invoking summarize_via_sampling, so no mock is needed for the error path. The test runs in milliseconds and never touches the sampling helper. If you find your test hangs or makes a real LLM call, your tool's early-return check on unknown ids is in the wrong order — fix the tool, not the test.

Solution to Task 3 — Redaction policy + caplog assertion

Part A — pkb/redaction_policy.md

# PKB Logging Redaction Policy

## Always safe to log
- Tool name, phase (`start` / `ok` / `err`), duration in ms.
- Exception **type** name (not the message).
- Trace id, parent trace id, request id.
- User id (the id itself — not the user's data).
- Counts and lengths: `n_results=3`, `text_len=120`, `result_type=str`.

## Must be redacted (truncate or summarise)
- Any string field longer than 200 characters — route through
  `observability.redact()` before logging.
- Any binary content — log only ``.

## Must NEVER appear in logs at any level
- Full tool arguments (may contain user prompts, file contents, secrets).
- Full LLM responses from sampling (user-paid output, both expensive to store
  and a leak risk if it echoed a prompt injection).
- `Authorization` headers, bearer tokens, API keys.
- Environment variable values (`PKB_USER`, `PKB_THIRD_PARTY_KEY`, anything
  from `os.environ`).
- `str(exception)` for exceptions raised by third-party SDKs — their
  messages routinely echo the failing URL with credentials in the query
  string. Log `type(e).__name__` only.

## Worked example
Bad:  `log.info("fetch ok body=%s", NOTES[id])`
Good: `log.info("tool=fetch phase=ok body_len=%d", len(NOTES[id]))`

The doc lives next to the code it constrains — inside the sample-server tree, not in the lesson HTML. A future contributor adding a new tool reads this before they decide what to log.

Part B — the matching caplog test

# Append to pkb/tests/test_tools.py

def test_fetch_does_not_log_body(caplog, fresh_notes):
    """The fetch tool must never write a note body into the captured log stream."""
    import asyncio
    with caplog.at_level("INFO", logger="pkb.tool.fetch"):
        async def go():
            async with create_connected_server_and_client_session(mcp._mcp_server) as client:
                return await client.call_tool("fetch", {"id": "hello"})
        asyncio.run(go())
    captured = " ".join(rec.getMessage() for rec in caplog.records)
    assert "Hello note body" not in captured

Why the test is structured this way:

• Scoped caplog.at_level(...). We pass logger="pkb.tool.fetch" so the captured records are only the ones emitted by the fetch tool's decorator-installed logger. Capturing at the root level invites flaky passes when test ordering rotates other loggers' records into the buffer.
• Sync test wrapping an async block. caplog's capture context is synchronous; the in-memory session is async. The asyncio.run(go()) pattern keeps the caplog context valid around the whole async lifetime. Marking this test with @pytest.mark.asyncio instead works too, but the sync wrapper is more obvious about what's being captured.
• Assertion uses joined record messages. " ".join(rec.getMessage() for rec in caplog.records) rebuilds the visible log lines deterministically. Using caplog.text works too; the joined form makes the failure message readable when the assertion does fire.

What you learned

• Trace propagation is a contract, not a feature. Once a composed tool exists, the parent id has to travel down or your log file becomes a stream of unrelated events — readable line-by-line, useless as a sequence.
• Parametrisation collapses N near-identical tests into one. The marginal cost of a new tool joining the "takes an id, can fail with unknown id" club is now one line in the parameter list, not a new test function.
• A written redaction policy plus one tripwire assertion is enough to catch the next mistake. The policy tells humans what to do; the assertion tells CI when they did it wrong. Both together cost about 50 lines and pay for themselves the first time someone tries to log str(e) for a third-party exception.

Module 5 (Plan E) leaves Python behind for the production half: TypeScript SDK, Docker, HTTP transport, multi-tenant remote servers, and production hardening — including the OpenTelemetry upgrade path for the tracing pattern you just built. Same shape; standardised wire format; distributed propagation across services.