Topic 5: Module 6 lab

"""Module 6 lab: a two-server host with namespacing, local rules, consent previews,
timeouts, no-progress detection, and graceful handling of a missing server."""
from __future__ import annotations

import json
import logging
import os
import shutil
import subprocess
import sys
import tempfile
import time
from pathlib import Path

import anyio

from mcp import Client

from m06_common import ROOT, ScriptedModel, estimate_tokens, notes_params
from m06_consent import make_preview_approve
from m06_failures import ResilientHost
from m06_no_progress import ProgressGuardHost
from m06_policy import PolicyHost, rule_for

logging.getLogger("notes_assistant.host").setLevel(logging.WARNING)
CALENDAR_URL = "http://127.0.0.1:8064/mcp"


class LabHost(ResilientHost, PolicyHost, ProgressGuardHost):
    """Resilient tool loading, then local rules and timeouts, then no-progress detection."""


def start_calendar_http() -> subprocess.Popen:
    env = os.environ | {"CALENDAR_TRANSPORT": "streamable-http", "CALENDAR_PORT": "8064", "PYTHONPATH": str(ROOT)}
    return subprocess.Popen([sys.executable, "examples/m06_calendar_server.py"], env=env, cwd=ROOT,
                            stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)


TURNS = [
    [("notes__search_notes", {"query": "nap study eeg"}),
     ("calendar__list_bookings", {"week_of": "2026-10-12", "resource": "eeg-caps"})],
    [("read_resource", {"uri": "notes://lab-sync-2026-09-02"})],
    [("calendar__book_slot", {"resource": "eeg-caps", "day": "2026-10-14", "booked_by": "Nare",
                              "note": "Nap study, second cohort"})],
    [("calendar__cancel_booking", {"booking_id": "b3"})],
    [("notes__create_note", {"title": "Nap study week 2", "tags": ["nap-study", "meeting"],
                             "body": "EEG caps booked for 14 October (Nare).\nSleep lab already booked by Tomas."})],
    "Booked the EEG caps for 14 October and saved [nap-study-week-2]. I could not cancel b3: your policy blocks it.",
]


async def main() -> None:
    calendar_server = start_calendar_http()
    prompts: list[str] = []
    try:
        await anyio.sleep(1.5)
        with tempfile.TemporaryDirectory() as tmp:
            notes_dir = Path(tmp) / "notes"
            shutil.copytree(ROOT / "notes", notes_dir)

            def answer(prompt: str) -> str:
                prompts.append(prompt)
                print(prompt + "y")
                return "y"

            async with Client(notes_params(notes_dir)) as notes, Client(CALENDAR_URL) as calendar:
                model = ScriptedModel(TURNS)
                host = LabHost({"notes": notes, "calendar": calendar}, chat_fn=model,
                               approve=make_preview_approve(notes_dir, answer=answer))
                started = time.perf_counter()
                result = await host.ask("Book the EEG caps for the nap study on 14 October, free up b3, and save a note.")
                seconds = time.perf_counter() - started
                specs = await host.load_tools()

            print(f"\nstop_reason={result.stop_reason}  rounds={model.rounds}  {seconds:.2f} s")
            print(f"{'call':<26} {'rule':<5} {'ok':<5} {'note':<17} ms")
            for c in result.calls:
                print(f"{c.name:<26} {rule_for(c.name) if c.name in host.tools else '-':<5} {c.ok!s:<5} {c.note or '-':<17} {c.duration_ms}")
            print(f"approval prompts: {len(prompts)} for {len(result.calls)} calls")
            print(f"tool definitions per round: {len(specs)} tools, ~{estimate_tokens(specs)[1]} tokens (estimate)")
            print(f"new files: {sorted(p.name for p in notes_dir.iterdir() if p.name.startswith('nap'))}")
            print(json.dumps({"answer": result.text}, indent=2))
    finally:
        calendar_server.kill()
        calendar_server.wait()


if __name__ == "__main__":
    anyio.run(main)