Tutorial D6: MCP Tooling & Architecture Basics

from datetime import datetime, timedelta

BASE = datetime(2025, 12, 24, 7, 0, 0)

def make_alarm(minutes: int, code: str, msg: str, line: str = "LINE-7"):
    return {
        "ts": (BASE + timedelta(minutes=minutes)).isoformat(timespec="seconds"),
        "line": line,
        "code": code,
        "msg": msg,
    }

ALARMS = [
    make_alarm(0, "PE_TIMEOUT", "Photo-eye blocked > 5s at infeed"),
    make_alarm(3, "VFD_OVERCURRENT", "Main conveyor VFD overcurrent trip"),
    make_alarm(6, "PE_TIMEOUT", "Photo-eye blocked > 5s at infeed"),
    make_alarm(9, "E_STOP", "E-stop circuit opened"),
    make_alarm(12, "PE_TIMEOUT", "Photo-eye blocked > 5s at infeed"),
]

# Agent 1: correct counting by full fault code
def agent1_parse_counts(alarms):
    counts = {}
    for a in alarms:
        counts[a["code"]] = counts.get(a["code"], 0) + 1
    return counts

# Agent 2: BUG — truncates fault code at underscore
def agent2_parse_counts(alarms):
    counts = {}
    for a in alarms:
        short = a["code"].split("_")[0]
        counts[short] = counts.get(short, 0) + 1
    return counts

print("Agent1 counts:", agent1_parse_counts(ALARMS))
print("Agent2 counts:", agent2_parse_counts(ALARMS))

from dataclasses import dataclass, field
from datetime import datetime, timezone
import json, hashlib

@dataclass
class ToolCall:
    """Records one tool invocation for audit purposes."""
    ts: str          # timestamp of call
    tool: str        # tool name
    args: dict       # arguments passed
    actor: str       # agent identifier
    ok: bool         # whether call succeeded
    reason: str      # "ok", "cache_hit", "unknown_tool", "write_blocked"

@dataclass
class ReadOnlyToolServer:
    """Minimal MCP-style tool host with read-only enforcement."""
    tools: dict = field(default_factory=dict)    # tool name → function
    audit: list = field(default_factory=list)    # ToolCall records
    cache: dict = field(default_factory=dict)    # cache key → result

def register(self, name: str, fn):
    """Register a tool function. Reject write_* tools by name."""
    if name.startswith("write_"):
        raise ValueError("Read-only server forbids write_* tools")
    self.tools[name] = fn

def list_tools(self):
    """Return sorted list of registered tool names."""
    return sorted(self.tools.keys())

# Attach methods onto the class so later cells can use srv.register(...)
ReadOnlyToolServer.register = register
ReadOnlyToolServer.list_tools = list_tools

def _cache_key(self, tool: str, args: dict) -> str:
    """Generate a deterministic cache key from tool name and arguments."""
    blob = json.dumps({"tool": tool, "args": args}, sort_keys=True).encode()
    return hashlib.sha256(blob).hexdigest()

ReadOnlyToolServer._cache_key = _cache_key

def call(self, tool: str, args: dict, actor: str):
    """Execute a tool call with security boundaries and audit trail."""
    now = datetime.now(timezone.utc).isoformat(timespec="seconds")

    # 1. Allowlist check: unknown tools are blocked
    if tool not in self.tools:
        self.audit.append(ToolCall(now, tool, args, actor, False, "unknown_tool"))
        raise KeyError(f"Unknown tool: {tool}")

    # 2. Read-only boundary: block write_* tools and suspicious args
    if tool.startswith("write_") or any(k.startswith("write") for k in args.keys()):
        self.audit.append(ToolCall(now, tool, args, actor, False, "write_blocked"))
        raise PermissionError("Read-only boundary: write operations blocked")

    # 3. Cache check: return cached result if available
    key = self._cache_key(tool, args)
    if key in self.cache:
        self.audit.append(ToolCall(now, tool, args, actor, True, "cache_hit"))
        return self.cache[key]

    # 4. Execute tool and cache result
    out = self.tools[tool](**args)
    self.cache[key] = out
    self.audit.append(ToolCall(now, tool, args, actor, True, "ok"))
    return out

ReadOnlyToolServer.call = call

# Reuse ALARMS + agent1_parse_counts from Experiment 1

def tool_alarm_counts(line: str):
    filtered = [a for a in ALARMS if a["line"] == line]
    return agent1_parse_counts(filtered)

srv = ReadOnlyToolServer()
srv.register("read_alarm_counts", tool_alarm_counts)

print("Tools:", srv.list_tools())

agent_a = "diagnostics_agent"
agent_b = "shift_report_agent"

res1 = srv.call("read_alarm_counts", {"line": "LINE-7"}, actor=agent_a)
res2 = srv.call("read_alarm_counts", {"line": "LINE-7"}, actor=agent_b)  # cache hit

print("A counts:", res1)
print("B counts:", res2)

# Unknown tool (blocked by allowlist)
try:
    srv.call("write_plc_tag", {"tag": "Conveyor.Start", "value": True}, actor=agent_a)
except Exception as e:
    print("Write attempt blocked:", type(e).__name__, str(e))

# Suspicious arg (blocked by write key detection)
try:
    srv.call("read_alarm_counts", {"line": "LINE-7", "write_override": True}, actor=agent_b)
except Exception as e:
    print("Arg write blocked:", type(e).__name__, str(e))

ok = sum(1 for c in srv.audit if c.ok)
blocked = sum(1 for c in srv.audit if (not c.ok))
cache_hits = sum(1 for c in srv.audit if c.reason == "cache_hit")

print(f"audit entries={len(srv.audit)} ok={ok} blocked={blocked} cache_hits={cache_hits}")
print("last 3 audit:")
for c in srv.audit[-3:]:
    print(c.tool, c.actor, c.ok, c.reason)