Homework — Lesson 10

Task 1 — Trace propagation through summarize_and_save

The log_tool_call decorator generates a fresh trace_id on every entry — great for a single tool call, useless for correlating an outer composed call with its inner sub-call. When a user invokes summarize_and_save, that tool internally awaits summarize_via_sampling; today the two emit unrelated trace ids and a log reader cannot tell which inner call belongs to which outer call.

Modify summarize_and_save in pkb/tools.py so that the trace_id generated for the composed call is threaded into the call to summarize_via_sampling, and update summarize_via_sampling so it accepts a parent_trace keyword and logs a line that includes it. The inner call still generates its own trace; it just additionally records its parent_trace.

What good output looks like, for a single call to summarize_and_save(id="hello"):

tool=summarize_and_save trace=AAAAAAAA phase=start
tool=summarize_via_sampling trace=BBBBBBBB parent_trace=AAAAAAAA phase=start
tool=summarize_via_sampling trace=BBBBBBBB parent_trace=AAAAAAAA phase=ok ...
tool=summarize_and_save trace=AAAAAAAA phase=ok ...

Hint: the cleanest place to generate the parent id is inside summarize_and_save itself (not inside the decorator). The decorator can keep doing its job; you just emit one additional manual log line at the top of the function body that carries the parent id, then pass that id explicitly into the helper.

✓ Self-check: call summarize_and_save(id="hello") in Inspector (with a mocked or real sampling client). The log shows one tool=summarize_and_save trace=AAA phase=start line, then tool=summarize_via_sampling trace=BBB parent_trace=AAA phase=start, then both finish in reverse order. AAA appears in exactly two lines; BBB in exactly two lines; AAA appears as parent_trace= exactly once.

Task 2 — Custom test pattern: parametrised error cases

The lab's test_tools.py contains test_fetch_unknown_id_is_error — one test, one tool. PKB actually has three tools that take an id and can fail with "unknown note id": fetch, summarize, and summarize_and_save. Collapse the boilerplate into a single @pytest.mark.parametrize'd test.

Create pkb/tests/test_errors.py with a single async test that:

• Takes parameters (tool_name, args) — one parameter row per affected tool, using a known-bad id like "no-such-note".
• Opens a connected in-memory session, calls the tool, and asserts both result.isError is True and that "unknown note id" appears in the first content block's text.
• Is decorated with both @pytest.mark.asyncio and @pytest.mark.parametrize (order: asyncio closest to the function, parametrize on top).

Bonus: summarize and summarize_and_save both call summarize_via_sampling internally — do you need to mock it to make the test pass? Think about the error path. If fetch's "unknown id" check fires before sampling is ever invoked, the mock is unnecessary for that case. Decide whether the same is true for the other two tools by reading their bodies in pkb/tools.py.

✓ Self-check: pytest pkb/tests/test_errors.py -v shows three PASSED lines, one per parameter case. Each parameter id in the pytest output identifies the tool by name (the parametrize ids default to the param values, which is the desired behaviour here).

Task 3 — Log-redaction policy (doc + assertion)

The lab's audit checklist lives in the lab guide. Promote it to a checked-in policy document next to the code it constrains, and back it up with an automated assertion so future edits can't quietly violate it.

Part A — Write the policy. Create mcp-servers/sample-server/python/pkb/redaction_policy.md with three sections (8–12 bullets total):

• Always safe to log — tool name, phase, duration, exception type, trace id, user id, counts, lengths.
• Must be redacted (truncate or summarise) — strings longer than 200 chars, binary content.
• Must NEVER appear in logs at any level — tool argument values (note bodies, search queries, the text field passed to echo), sampling responses, authorization headers, env-var values, stringified third-party exceptions.

Close the document with a small worked example: a "before" log line that leaks a note body, and the redacted "after" line that logs body_len= instead.

Part B — Prove it with a test. Add one new assertion to pkb/tests/test_tools.py — a test named test_fetch_does_not_log_body — that uses pytest's caplog fixture to capture log records emitted during a fetch(id="hello") call, then asserts that the literal seed text "Hello note body" does not appear anywhere in the captured log output.

Two implementation notes worth thinking about before you write the code:

• Use caplog.at_level("INFO", logger="pkb.tool.fetch") inside the test body so the assertion is scoped to the logger you actually care about. Asserting against caplog.text at the root level invites flaky passes when test ordering changes.
• The test is sync (no @pytest.mark.asyncio) but the in-memory session is async. Wrap the body in a nested async def and run it with asyncio.run(...) — this keeps caplog's synchronous capture context valid around the entire async block.

✓ Self-check: cat pkb/redaction_policy.md prints the policy with the three sections and worked example. pytest pkb/tests/test_tools.py::test_fetch_does_not_log_body -q exits with 1 passed. As an extra spike: temporarily add a log.info("fetch body=%s", NOTES["hello"]) line to fetch and verify the new test fails — that's the proof that the assertion would catch a future regression.