From a573142d108fda6f621b8873813b1797f44cac54 Mon Sep 17 00:00:00 2001 From: chrisbartoloburlo Date: Thu, 23 Apr 2026 12:31:30 +0200 Subject: [PATCH 1/2] Add Protocol Builder tab for session type protocol design Adds a new "Protocol Builder" tab to the MCP Inspector that lets users visually compose interaction protocols from discovered MCP tools using session type theory. Features: - Discovers tools via tools/list and presents them as clickable blocks - Supports send (!), receive (?), internal/external choice, and recursion - Live DSL output with syntax highlighting - State machine visualization - Generates Python integration code using llmcontract - Export/download protocol as Python file This enables runtime protocol monitoring for LLM agent interactions, ensuring tools are called in the correct order per a defined contract. Co-Authored-By: Claude Opus 4.6 --- client/src/App.tsx | 14 + client/src/components/ProtocolBuilderTab.tsx | 913 +++++++++++++++++++ 2 files changed, 927 insertions(+) create mode 100644 client/src/components/ProtocolBuilderTab.tsx diff --git a/client/src/App.tsx b/client/src/App.tsx index 59d15ba06..ac9af76bf 100644 --- a/client/src/App.tsx +++ b/client/src/App.tsx @@ -63,6 +63,7 @@ import { ListTodo, MessageSquare, Settings, + Shield, } from "lucide-react"; import { z } from "zod"; @@ -79,6 +80,7 @@ import Sidebar from "./components/Sidebar"; import ToolsTab from "./components/ToolsTab"; import TasksTab from "./components/TasksTab"; import AppsTab from "./components/AppsTab"; +import ProtocolBuilderTab from "./components/ProtocolBuilderTab"; import { InspectorConfig } from "./lib/configurationTypes"; import { getMCPProxyAddress, @@ -346,6 +348,7 @@ const App = () => { "roots", "auth", "metadata", + "protocol-builder", ]; if (!validTabs.includes(originatingTab)) return; @@ -1437,6 +1440,10 @@ const App = () => { Metadata + + + Protocol Builder +
@@ -1673,6 +1680,13 @@ const App = () => { metadata={metadata} onMetadataChange={handleMetadataChange} /> + { + clearError("tools"); + listTools(); + }} + /> )}
diff --git a/client/src/components/ProtocolBuilderTab.tsx b/client/src/components/ProtocolBuilderTab.tsx new file mode 100644 index 000000000..bee756690 --- /dev/null +++ b/client/src/components/ProtocolBuilderTab.tsx @@ -0,0 +1,913 @@ +import { useState, useCallback, useMemo } from "react"; +import { TabsContent } from "@/components/ui/tabs"; +import { Button } from "@/components/ui/button"; +import { Tool } from "@modelcontextprotocol/sdk/types.js"; +import { + GripVertical, + Plus, + Trash2, + Copy, + RotateCcw, + ChevronDown, + ChevronRight, + Download, +} from "lucide-react"; + +// ── Types ────────────────────────────────────────────────── + +type Direction = "send" | "receive"; + +interface ProtocolStep { + id: string; + type: "action" | "choice" | "recursion"; + direction?: Direction; + label?: string; + toolName?: string; + branches?: ProtocolBranch[]; + recVar?: string; + isRecRef?: boolean; +} + +interface ProtocolBranch { + id: string; + label: string; + steps: ProtocolStep[]; +} + +// ── Helpers ──────────────────────────────────────────────── + +let nextId = 0; +const uid = () => `step-${++nextId}`; + +function stepsToProtocol(steps: ProtocolStep[]): string { + if (steps.length === 0) return "end"; + + const parts: string[] = []; + + for (const step of steps) { + if (step.type === "action") { + if (step.isRecRef && step.recVar) { + parts.push(step.recVar); + } else { + const prefix = step.direction === "send" ? "!" : "?"; + parts.push(`${prefix}${step.label}`); + } + } else if (step.type === "choice" && step.branches) { + const prefix = step.direction === "send" ? "!" : "?"; + const branchStrs = step.branches.map((b) => { + const inner = stepsToProtocol(b.steps); + return inner === "end" ? b.label : `${b.label}.${inner}`; + }); + parts.push(`${prefix}{${branchStrs.join(", ")}}`); + } else if (step.type === "recursion" && step.recVar) { + const inner = stepsToProtocol(steps.slice(steps.indexOf(step) + 1)); + const body = + step.branches && step.branches.length > 0 + ? stepsToProtocol( + step.branches[0].steps.concat( + steps.slice(steps.indexOf(step) + 1), + ), + ) + : inner; + return `rec ${step.recVar}.${body}`; + } + } + + if (parts.length === 0) return "end"; + return parts.join(".") + ".end"; +} + +function generatePythonSnippet(protocol: string, tools: Tool[]): string { + const toolNames = tools.map((t) => t.name).slice(0, 5); + const toolEntries = toolNames + .map((n) => ` "${n}": ${n}_fn,`) + .join("\n"); + + return `from llmcontract import Monitor, MonitoredClient, ToolMiddleware, LLMResponse + +# Define the protocol +protocol = "${protocol}" + +# Create a shared monitor +monitor = Monitor(protocol) + +# Wrap your LLM client +client = MonitoredClient( + llm_call=your_llm_fn, + response_adapter=your_adapter, + monitor=monitor, + send_label="Request", + receive_label=lambda r: "ToolCall" if r.has_tool_calls else "FinalAnswer", +) + +# Register tools with the middleware +tools = ToolMiddleware( + monitor=monitor, + tools={ +${toolEntries || ' # "tool_name": tool_fn,'} + }, +)`; +} + +// ── Protocol Builder Component ───────────────────────────── + +const ProtocolBuilderTab = ({ + tools, + listTools, +}: { + tools: Tool[]; + listTools: () => void; +}) => { + const [steps, setSteps] = useState([]); + const [expandedTools, setExpandedTools] = useState(true); + const [copied, setCopied] = useState<"dsl" | "python" | null>(null); + + const protocol = useMemo(() => { + if (steps.length === 0) return "end"; + return stepsToProtocol(steps); + }, [steps]); + + const pythonSnippet = useMemo( + () => generatePythonSnippet(protocol, tools), + [protocol, tools], + ); + + const addStep = useCallback((step: ProtocolStep) => { + setSteps((prev) => [...prev, step]); + }, []); + + const removeStep = useCallback((id: string) => { + setSteps((prev) => prev.filter((s) => s.id !== id)); + }, []); + + const addToolAsSteps = useCallback((tool: Tool) => { + const sendStep: ProtocolStep = { + id: uid(), + type: "action", + direction: "send", + label: tool.name, + toolName: tool.name, + }; + const recvStep: ProtocolStep = { + id: uid(), + type: "action", + direction: "receive", + label: `${tool.name}Result`, + toolName: tool.name, + }; + setSteps((prev) => [...prev, sendStep, recvStep]); + }, []); + + const addChoice = useCallback((direction: Direction) => { + const step: ProtocolStep = { + id: uid(), + type: "choice", + direction, + branches: [ + { id: uid(), label: "BranchA", steps: [] }, + { id: uid(), label: "BranchB", steps: [] }, + ], + }; + setSteps((prev) => [...prev, step]); + }, []); + + const addRecursion = useCallback(() => { + const varName = `X${steps.filter((s) => s.type === "recursion").length}`; + const recStep: ProtocolStep = { + id: uid(), + type: "recursion", + recVar: varName, + }; + setSteps((prev) => [...prev, recStep]); + }, [steps]); + + const addRecRef = useCallback((varName: string) => { + const step: ProtocolStep = { + id: uid(), + type: "action", + isRecRef: true, + recVar: varName, + }; + setSteps((prev) => [...prev, step]); + }, []); + + const updateBranchLabel = useCallback( + (stepId: string, branchId: string, label: string) => { + setSteps((prev) => + prev.map((s) => { + if (s.id !== stepId || !s.branches) return s; + return { + ...s, + branches: s.branches.map((b) => + b.id === branchId ? { ...b, label } : b, + ), + }; + }), + ); + }, + [], + ); + + const addBranch = useCallback((stepId: string) => { + setSteps((prev) => + prev.map((s) => { + if (s.id !== stepId || !s.branches) return s; + return { + ...s, + branches: [ + ...s.branches, + { + id: uid(), + label: `Branch${String.fromCharCode(65 + s.branches.length)}`, + steps: [], + }, + ], + }; + }), + ); + }, []); + + const removeBranch = useCallback((stepId: string, branchId: string) => { + setSteps((prev) => + prev.map((s) => { + if (s.id !== stepId || !s.branches) return s; + if (s.branches.length <= 2) return s; + return { + ...s, + branches: s.branches.filter((b) => b.id !== branchId), + }; + }), + ); + }, []); + + const copyToClipboard = useCallback( + (text: string, type: "dsl" | "python") => { + navigator.clipboard.writeText(text); + setCopied(type); + setTimeout(() => setCopied(null), 2000); + }, + [], + ); + + const clearAll = useCallback(() => { + setSteps([]); + }, []); + + const recVars = steps + .filter((s) => s.type === "recursion" && s.recVar) + .map((s) => s.recVar!); + + return ( + +
+
+ {/* ── Left: Tool Palette ── */} +
+
+

+ MCP Tools +

+ +
+ + {tools.length === 0 ? ( +
+

+ No tools discovered yet +

+ +
+ ) : ( +
+ + {expandedTools && + tools.map((tool) => ( + + ))} +
+ )} + +
+

+ Protocol Constructs +

+ + + + + + {recVars.length > 0 && ( +
+ {recVars.map((v) => ( + + ))} +
+ )} +
+
+ + {/* ── Middle: Protocol Steps ── */} +
+
+

+ Protocol Sequence +

+ +
+ + {steps.length === 0 ? ( +
+

+ Click tools or constructs on the left to build your protocol +

+
+ ) : ( +
+ {steps.map((step) => ( + removeStep(step.id)} + onUpdateLabel={(label) => + setSteps((prev) => + prev.map((s) => + s.id === step.id ? { ...s, label } : s, + ), + ) + } + onUpdateBranchLabel={updateBranchLabel} + onAddBranch={() => addBranch(step.id)} + onRemoveBranch={(branchId) => + removeBranch(step.id, branchId) + } + /> + ))} + + {/* Terminal indicator */} +
+ + end + +
+
+ )} +
+ + {/* ── Right: Output ── */} +
+

+ Output +

+ + {/* DSL */} +
+
+

+ Session Type DSL +

+ +
+ 
+                
+
+
+ + {/* State machine preview */} +
+

+ State Machine Preview +

+
+ +
+
+ + {/* Python snippet */} +
+
+

+ Python Integration +

+ +
+ 
+                {pythonSnippet}
+
+
+ + {/* Download button */} + +
+
+
+ + ); +}; + +// ── Step Card ────────────────────────────────────────────── + +function StepCard({ + step, + onRemove, + onUpdateLabel, + onUpdateBranchLabel, + onAddBranch, + onRemoveBranch, +}: { + step: ProtocolStep; + onRemove: () => void; + onUpdateLabel: (label: string) => void; + onUpdateBranchLabel: ( + stepId: string, + branchId: string, + label: string, + ) => void; + onAddBranch: () => void; + onRemoveBranch: (branchId: string) => void; +}) { + if (step.type === "action" && step.isRecRef) { + return ( +
+ + + ↻ loop → {step.recVar} + +
+ +
+ ); + } + + if (step.type === "action") { + const isSend = step.direction === "send"; + const bgClass = isSend + ? "bg-green-50 dark:bg-green-950 border-green-200 dark:border-green-800" + : "bg-blue-50 dark:bg-blue-950 border-blue-200 dark:border-blue-800"; + const prefixClass = isSend + ? "text-green-600 dark:text-green-400" + : "text-blue-600 dark:text-blue-400"; + + return ( +
+ + + {isSend ? "!" : "?"} + + onUpdateLabel(e.target.value)} + className="font-mono text-xs bg-transparent border-none outline-none flex-1 min-w-0" + placeholder="label" + /> + {step.toolName && ( + + ({step.toolName}) + + )} + +
+ ); + } + + if (step.type === "choice") { + const isSend = step.direction === "send"; + const borderClass = isSend + ? "border-green-300 dark:border-green-700" + : "border-blue-300 dark:border-blue-700"; + const prefixClass = isSend + ? "text-green-600 dark:text-green-400" + : "text-blue-600 dark:text-blue-400"; + + return ( +
+
+ + + {isSend ? "!" : "?"} + {"{"} + + + {isSend ? "Internal" : "External"} Choice + +
+ + +
+
+ {step.branches?.map((branch) => ( +
+ + onUpdateBranchLabel(step.id, branch.id, e.target.value) + } + className="font-mono text-xs bg-transparent border-none outline-none flex-1 min-w-0" + placeholder="branch label" + /> + {(step.branches?.length ?? 0) > 2 && ( + + )} +
+ ))} +
+
+ + {"}"} + +
+
+ ); + } + + if (step.type === "recursion") { + return ( +
+ + + rec + + + {step.recVar} + + . +
+ +
+ ); + } + + return null; +} + +// ── Protocol Syntax Highlighting ─────────────────────────── + +function ProtocolHighlight({ protocol }: { protocol: string }) { + const tokens: { text: string; cls: string }[] = []; + let i = 0; + const src = protocol; + + while (i < src.length) { + if (src[i] === "!") { + tokens.push({ text: "!", cls: "text-green-600 dark:text-green-400" }); + i++; + } else if (src[i] === "?") { + tokens.push({ text: "?", cls: "text-blue-600 dark:text-blue-400" }); + i++; + } else if ( + src[i] === "." || + src[i] === "{" || + src[i] === "}" || + src[i] === "," + ) { + tokens.push({ text: src[i], cls: "text-gray-400" }); + i++; + } else if ( + src.slice(i, i + 3) === "rec" && + (i + 3 >= src.length || !/[a-zA-Z0-9_]/.test(src[i + 3])) + ) { + tokens.push({ + text: "rec", + cls: "text-amber-600 dark:text-amber-400 italic", + }); + i += 3; + } else if ( + src.slice(i, i + 3) === "end" && + (i + 3 >= src.length || !/[a-zA-Z0-9_]/.test(src[i + 3])) + ) { + tokens.push({ text: "end", cls: "text-gray-500 italic" }); + i += 3; + } else if (/[a-zA-Z0-9_]/.test(src[i])) { + const start = i; + while (i < src.length && /[a-zA-Z0-9_]/.test(src[i])) i++; + tokens.push({ + text: src.slice(start, i), + cls: "text-gray-900 dark:text-gray-100", + }); + } else if (src[i] === " ") { + tokens.push({ text: " ", cls: "" }); + i++; + } else { + tokens.push({ text: src[i], cls: "" }); + i++; + } + } + + return ( + <> + {tokens.map((t, idx) => ( + + {t.text} + + ))} + + ); +} + +// ── State Machine Preview ────────────────────────────────── + +function StateMachinePreview({ protocol }: { protocol: string }) { + // Simple parser → FSM for visualization + const { states, transitions } = useMemo(() => { + try { + return parseToFSM(protocol); + } catch { + return { states: [], transitions: [] }; + } + }, [protocol]); + + if (states.length === 0) { + return ( +

+ Add steps to see the state machine +

+ ); + } + + return ( +
+ {transitions.map((t, i) => ( +
+ + S{t.from} + + → + + {t.dir === "send" ? "!" : t.dir === "receive" ? "?" : ""} + {t.label} + + → +
+ ))} + + S{states.length > 0 ? states[states.length - 1] : 0} + +
+ ); +} + +function parseToFSM(protocol: string): { + states: number[]; + transitions: { from: number; to: number; dir: string; label: string }[]; +} { + const states: number[] = [0]; + const transitions: { + from: number; + to: number; + dir: string; + label: string; + }[] = []; + let stateId = 0; + let i = 0; + const src = protocol; + + const skipWS = () => { + while (i < src.length && " \t\n\r".includes(src[i])) i++; + }; + + const readIdent = () => { + skipWS(); + const start = i; + while (i < src.length && /[a-zA-Z0-9_]/.test(src[i])) i++; + return src.slice(start, i); + }; + + const parseSimple = () => { + skipWS(); + if (i >= src.length) return; + + if (src[i] === "!" || src[i] === "?") { + const dir = src[i] === "!" ? "send" : "receive"; + i++; + skipWS(); + if (i < src.length && src[i] === "{") { + // choice + i++; + while (i < src.length && src[i] !== "}") { + skipWS(); + const label = readIdent(); + if (label) { + const from = stateId; + stateId++; + states.push(stateId); + transitions.push({ from, to: stateId, dir, label }); + } + skipWS(); + if (src[i] === ".") { + i++; + parseSimple(); + } + if (src[i] === ",") i++; + } + if (src[i] === "}") i++; + } else { + const label = readIdent(); + if (label) { + const from = stateId; + stateId++; + states.push(stateId); + transitions.push({ from, to: stateId, dir, label }); + } + } + } else if (src.slice(i, i + 3) === "rec") { + i += 3; + readIdent(); // var name + skipWS(); + if (src[i] === ".") i++; + parseSimple(); + return; + } else if (src.slice(i, i + 3) === "end") { + i += 3; + return; + } else { + // recvar or label + const ident = readIdent(); + if (ident && ident !== "end") { + // rec var reference — skip + } + } + + skipWS(); + if (src[i] === ".") { + i++; + parseSimple(); + } + }; + + try { + parseSimple(); + } catch { + // graceful fallback + } + + return { states, transitions }; +} + +export default ProtocolBuilderTab; From 94a9c7be5dbb13071210d5c10fe70843f5a52758 Mon Sep 17 00:00:00 2001 From: chrisbartoloburlo Date: Thu, 23 Apr 2026 13:56:49 +0200 Subject: [PATCH 2/2] Enhance Protocol Builder with convert-to-choice, paired steps, and FSM preview MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add convert-to-choice buttons on PairCard: convert ! to !{} or ? to ?{} inline, keeping the partner as a standalone step - ? to ?{} auto-populates branches with {sendLabel}Result and {sendLabel}Error - ! to !{} populates branches with available MCP tool names - Rewrite parseToFSM to correctly model choices as branching from same state - Update StateMachinePreview to show branching visually with grouped transitions - Fix recursion (rec X.) not being terminal — only loop-back refs terminate - Support hyphens in DSL identifiers for MCP tool name compatibility - Fix nested convert-to-choice crash (wrong variable name in recursive call) - Fix package name: llmcontract -> llmsessioncontract - Add example agent using MCP server-everything tools (echo, get-sum) - Register protocol-builder tab in all validTabs arrays Co-Authored-By: Claude Opus 4.6 --- client/src/App.tsx | 2 + client/src/components/ProtocolBuilderTab.tsx | 1485 ++++++++++++++---- examples/weather_agent.py | 108 ++ 3 files changed, 1320 insertions(+), 275 deletions(-) create mode 100644 examples/weather_agent.py diff --git a/client/src/App.tsx b/client/src/App.tsx index ac9af76bf..497e29c3c 100644 --- a/client/src/App.tsx +++ b/client/src/App.tsx @@ -483,6 +483,7 @@ const App = () => { "roots", "auth", "metadata", + "protocol-builder", ]; const isValidTab = validTabs.includes(hash); @@ -817,6 +818,7 @@ const App = () => { "roots", "auth", "metadata", + "protocol-builder", ]; if (validTabs.includes(originatingTab)) { diff --git a/client/src/components/ProtocolBuilderTab.tsx b/client/src/components/ProtocolBuilderTab.tsx index bee756690..8f242d334 100644 --- a/client/src/components/ProtocolBuilderTab.tsx +++ b/client/src/components/ProtocolBuilderTab.tsx @@ -11,6 +11,8 @@ import { ChevronDown, ChevronRight, Download, + Target, + GitBranch, } from "lucide-react"; // ── Types ────────────────────────────────────────────────── @@ -26,6 +28,7 @@ interface ProtocolStep { branches?: ProtocolBranch[]; recVar?: string; isRecRef?: boolean; + pairId?: string; // links ! and ? steps together } interface ProtocolBranch { @@ -34,6 +37,12 @@ interface ProtocolBranch { steps: ProtocolStep[]; } +// Target for where palette clicks should add steps +interface InsertTarget { + choiceStepId: string; + branchId: string; +} + // ── Helpers ──────────────────────────────────────────────── let nextId = 0; @@ -83,7 +92,7 @@ function generatePythonSnippet(protocol: string, tools: Tool[]): string { .map((n) => ` "${n}": ${n}_fn,`) .join("\n"); - return `from llmcontract import Monitor, MonitoredClient, ToolMiddleware, LLMResponse + return `from llmsessioncontract import Monitor, MonitoredClient, ToolMiddleware, LLMResponse # Define the protocol protocol = "${protocol}" @@ -109,6 +118,209 @@ ${toolEntries || ' # "tool_name": tool_fn,'} )`; } +// ── Recursive state helpers ──────────────────────────────── + +function updateStepDeep( + steps: ProtocolStep[], + stepId: string, + updater: (step: ProtocolStep) => ProtocolStep, +): ProtocolStep[] { + return steps.map((s) => { + if (s.id === stepId) return updater(s); + if (s.branches) { + return { + ...s, + branches: s.branches.map((b) => ({ + ...b, + steps: updateStepDeep(b.steps, stepId, updater), + })), + }; + } + return s; + }); +} + +function removeStepDeep(steps: ProtocolStep[], stepId: string): ProtocolStep[] { + // Find the step to get its pairId + const pairId = findPairId(steps, stepId); + const shouldRemove = (s: ProtocolStep) => + s.id === stepId || (pairId && s.pairId === pairId); + + const filtered = steps.filter((s) => !shouldRemove(s)); + return filtered.map((s) => { + if (s.branches) { + return { + ...s, + branches: s.branches.map((b) => ({ + ...b, + steps: removeStepDeep(b.steps, stepId), + })), + }; + } + return s; + }); +} + +function findPairId(steps: ProtocolStep[], stepId: string): string | undefined { + for (const s of steps) { + if (s.id === stepId) return s.pairId; + if (s.branches) { + for (const b of s.branches) { + const found = findPairId(b.steps, stepId); + if (found) return found; + } + } + } + return undefined; +} + +function addStepToBranch( + steps: ProtocolStep[], + choiceStepId: string, + branchId: string, + newStep: ProtocolStep, +): ProtocolStep[] { + return steps.map((s) => { + if (s.id === choiceStepId && s.branches) { + return { + ...s, + branches: s.branches.map((b) => + b.id === branchId ? { ...b, steps: [...b.steps, newStep] } : b, + ), + }; + } + if (s.branches) { + return { + ...s, + branches: s.branches.map((b) => ({ + ...b, + steps: addStepToBranch(b.steps, choiceStepId, branchId, newStep), + })), + }; + } + return s; + }); +} + +function isTerminated(steps: ProtocolStep[]): boolean { + if (steps.length === 0) return false; + const last = steps[steps.length - 1]; + if (last.type === "choice") return true; + // Note: "recursion" (rec X.) is NOT terminal — it's a scope opener. + // Only a rec *reference* (loop back to X) is terminal. + if (last.type === "action" && last.isRecRef) return true; + return false; +} + +function collectRecVars(steps: ProtocolStep[]): string[] { + const vars: string[] = []; + for (const s of steps) { + if (s.type === "recursion" && s.recVar) vars.push(s.recVar); + if (s.branches) { + for (const b of s.branches) { + vars.push(...collectRecVars(b.steps)); + } + } + } + return vars; +} + +function collectSendLabels(steps: ProtocolStep[]): string[] { + const labels: string[] = []; + for (const s of steps) { + if ( + s.type === "action" && + s.direction === "send" && + s.label && + !s.isRecRef + ) { + labels.push(s.label); + } + if (s.branches) { + for (const b of s.branches) { + labels.push(...collectSendLabels(b.steps)); + } + } + } + return [...new Set(labels)]; +} + +function deriveReceiveOptions(sendLabels: string[]): string[] { + const options: string[] = []; + for (const label of sendLabels) { + options.push(`${label}Result`); + options.push(`${label}Error`); + } + return options; +} + +// Check if a branch target is terminated +function isBranchTerminated( + steps: ProtocolStep[], + choiceStepId: string, + branchId: string, +): boolean { + for (const s of steps) { + if (s.id === choiceStepId && s.branches) { + const branch = s.branches.find((b) => b.id === branchId); + return branch ? isTerminated(branch.steps) : false; + } + if (s.branches) { + for (const b of s.branches) { + const result = isBranchTerminated(b.steps, choiceStepId, branchId); + if (result) return true; + } + } + } + return false; +} + +// Replace a paired action with a choice, keeping its partner as an unpaired action +function convertPairToChoice( + steps: ProtocolStep[], + stepId: string, + pairId: string, + direction: Direction, + branchLabels: string[], +): ProtocolStep[] { + return steps.map((s) => { + // The step being converted → becomes a choice + if (s.id === stepId) { + return { + id: s.id, + type: "choice" as const, + direction, + branches: branchLabels.map((label) => ({ + id: uid(), + label, + steps: [], + })), + }; + } + // The partner step → just remove the pairId so it becomes standalone + if (s.pairId === pairId && s.id !== stepId) { + const { pairId: _, ...rest } = s; + return rest; + } + if (s.branches) { + return { + ...s, + branches: s.branches.map((b) => ({ + ...b, + steps: convertPairToChoice( + b.steps, + stepId, + pairId, + direction, + branchLabels, + ), + })), + }; + } + return s; + }); +} + // ── Protocol Builder Component ───────────────────────────── const ProtocolBuilderTab = ({ @@ -121,6 +333,7 @@ const ProtocolBuilderTab = ({ const [steps, setSteps] = useState([]); const [expandedTools, setExpandedTools] = useState(true); const [copied, setCopied] = useState<"dsl" | "python" | null>(null); + const [insertTarget, setInsertTarget] = useState(null); const protocol = useMemo(() => { if (steps.length === 0) return "end"; @@ -132,113 +345,159 @@ const ProtocolBuilderTab = ({ [protocol, tools], ); - const addStep = useCallback((step: ProtocolStep) => { - setSteps((prev) => [...prev, step]); - }, []); - - const removeStep = useCallback((id: string) => { - setSteps((prev) => prev.filter((s) => s.id !== id)); - }, []); + const topTerminated = isTerminated(steps); + + // Where should palette clicks go? + const targetTerminated = insertTarget + ? isBranchTerminated( + steps, + insertTarget.choiceStepId, + insertTarget.branchId, + ) + : topTerminated; + + const addStepToTarget = useCallback( + (step: ProtocolStep) => { + if (insertTarget) { + setSteps((prev) => + addStepToBranch( + prev, + insertTarget.choiceStepId, + insertTarget.branchId, + step, + ), + ); + } else { + setSteps((prev) => [...prev, step]); + } + }, + [insertTarget], + ); - const addToolAsSteps = useCallback((tool: Tool) => { - const sendStep: ProtocolStep = { - id: uid(), - type: "action", - direction: "send", - label: tool.name, - toolName: tool.name, - }; - const recvStep: ProtocolStep = { - id: uid(), - type: "action", - direction: "receive", - label: `${tool.name}Result`, - toolName: tool.name, - }; - setSteps((prev) => [...prev, sendStep, recvStep]); - }, []); + const addToolAsSteps = useCallback( + (tool: Tool) => { + const pair = uid(); + const sendStep: ProtocolStep = { + id: uid(), + type: "action", + direction: "send", + label: tool.name, + toolName: tool.name, + pairId: pair, + }; + const recvStep: ProtocolStep = { + id: uid(), + type: "action", + direction: "receive", + label: `${tool.name}Result`, + toolName: tool.name, + pairId: pair, + }; + if (insertTarget) { + setSteps((prev) => { + let s = addStepToBranch( + prev, + insertTarget.choiceStepId, + insertTarget.branchId, + sendStep, + ); + s = addStepToBranch( + s, + insertTarget.choiceStepId, + insertTarget.branchId, + recvStep, + ); + return s; + }); + } else { + setSteps((prev) => [...prev, sendStep, recvStep]); + } + }, + [insertTarget], + ); - const addChoice = useCallback((direction: Direction) => { - const step: ProtocolStep = { - id: uid(), - type: "choice", - direction, - branches: [ - { id: uid(), label: "BranchA", steps: [] }, - { id: uid(), label: "BranchB", steps: [] }, - ], - }; - setSteps((prev) => [...prev, step]); - }, []); + const addChoice = useCallback( + (direction: Direction) => { + const step: ProtocolStep = { + id: uid(), + type: "choice", + direction, + branches: [ + { id: uid(), label: "BranchA", steps: [] }, + { id: uid(), label: "BranchB", steps: [] }, + ], + }; + addStepToTarget(step); + }, + [addStepToTarget], + ); const addRecursion = useCallback(() => { - const varName = `X${steps.filter((s) => s.type === "recursion").length}`; - const recStep: ProtocolStep = { - id: uid(), - type: "recursion", - recVar: varName, - }; - setSteps((prev) => [...prev, recStep]); - }, [steps]); + setSteps((prev) => { + const varName = `X${collectRecVars(prev).length}`; + const recStep: ProtocolStep = { + id: uid(), + type: "recursion", + recVar: varName, + }; + if (insertTarget) { + return addStepToBranch( + prev, + insertTarget.choiceStepId, + insertTarget.branchId, + recStep, + ); + } + return [...prev, recStep]; + }); + }, [insertTarget]); + + const addRecRef = useCallback( + (varName: string) => { + const step: ProtocolStep = { + id: uid(), + type: "action", + isRecRef: true, + recVar: varName, + }; + addStepToTarget(step); + }, + [addStepToTarget], + ); - const addRecRef = useCallback((varName: string) => { - const step: ProtocolStep = { - id: uid(), - type: "action", - isRecRef: true, - recVar: varName, - }; - setSteps((prev) => [...prev, step]); + const handleUpdateStep = useCallback( + (stepId: string, updater: (step: ProtocolStep) => ProtocolStep) => { + setSteps((prev) => updateStepDeep(prev, stepId, updater)); + }, + [], + ); + + const handleRemoveStep = useCallback((stepId: string) => { + setSteps((prev) => removeStepDeep(prev, stepId)); }, []); - const updateBranchLabel = useCallback( - (stepId: string, branchId: string, label: string) => { + const handleConvertToChoice = useCallback( + ( + stepId: string, + pairId: string, + direction: Direction, + branchLabels: string[], + ) => { setSteps((prev) => - prev.map((s) => { - if (s.id !== stepId || !s.branches) return s; - return { - ...s, - branches: s.branches.map((b) => - b.id === branchId ? { ...b, label } : b, - ), - }; - }), + convertPairToChoice(prev, stepId, pairId, direction, branchLabels), ); }, [], ); - const addBranch = useCallback((stepId: string) => { - setSteps((prev) => - prev.map((s) => { - if (s.id !== stepId || !s.branches) return s; - return { - ...s, - branches: [ - ...s.branches, - { - id: uid(), - label: `Branch${String.fromCharCode(65 + s.branches.length)}`, - steps: [], - }, - ], - }; - }), - ); - }, []); - - const removeBranch = useCallback((stepId: string, branchId: string) => { - setSteps((prev) => - prev.map((s) => { - if (s.id !== stepId || !s.branches) return s; - if (s.branches.length <= 2) return s; - return { - ...s, - branches: s.branches.filter((b) => b.id !== branchId), - }; - }), - ); - }, []); + const handleAddStepToBranch = useCallback( + (choiceStepId: string, branchId: string, newStep: ProtocolStep) => { + setSteps((prev) => + addStepToBranch(prev, choiceStepId, branchId, newStep), + ); + }, + [], + ); const copyToClipboard = useCallback( (text: string, type: "dsl" | "python") => { @@ -251,11 +510,33 @@ const ProtocolBuilderTab = ({ const clearAll = useCallback(() => { setSteps([]); + setInsertTarget(null); }, []); - const recVars = steps - .filter((s) => s.type === "recursion" && s.recVar) - .map((s) => s.recVar!); + const recVars = collectRecVars(steps); + const sendLabels = collectSendLabels(steps); + const receiveOptions = deriveReceiveOptions(sendLabels); + + // Find the target branch label for display + const targetLabel = useMemo(() => { + if (!insertTarget) return null; + const findBranch = (steps: ProtocolStep[]): string | null => { + for (const s of steps) { + if (s.id === insertTarget.choiceStepId && s.branches) { + const b = s.branches.find((b) => b.id === insertTarget.branchId); + return b ? b.label : null; + } + if (s.branches) { + for (const b of s.branches) { + const found = findBranch(b.steps); + if (found) return found; + } + } + } + return null; + }; + return findBranch(steps); + }, [insertTarget, steps]); return ( @@ -272,6 +553,22 @@ const ProtocolBuilderTab = ({
+ {/* Insert target indicator */} + {insertTarget && ( +
+ + + Adding to: {targetLabel} + + +
+ )} + {tools.length === 0 ? (

@@ -298,8 +595,8 @@ const ProtocolBuilderTab = ({ tools.map((tool) => ( -

) : (
- {steps.map((step) => ( - removeStep(step.id)} - onUpdateLabel={(label) => - setSteps((prev) => - prev.map((s) => - s.id === step.id ? { ...s, label } : s, - ), - ) - } - onUpdateBranchLabel={updateBranchLabel} - onAddBranch={() => addBranch(step.id)} - onRemoveBranch={(branchId) => - removeBranch(step.id, branchId) - } - /> - ))} + {/* Terminal indicator */}
@@ -536,26 +829,299 @@ const ProtocolBuilderTab = ({ ); }; +// ── Step List (groups paired !? steps) ───────────────────── + +function StepList({ + steps, + tools, + recVars, + receiveOptions, + insertTarget, + onSetInsertTarget, + onUpdateStep, + onRemoveStep, + onAddStepToBranch, + onConvertToChoice, + depth = 0, +}: { + steps: ProtocolStep[]; + tools: Tool[]; + recVars: string[]; + receiveOptions: string[]; + insertTarget: InsertTarget | null; + onSetInsertTarget: (target: InsertTarget | null) => void; + onUpdateStep: ( + stepId: string, + updater: (s: ProtocolStep) => ProtocolStep, + ) => void; + onRemoveStep: (stepId: string) => void; + onAddStepToBranch: ( + choiceStepId: string, + branchId: string, + newStep: ProtocolStep, + ) => void; + onConvertToChoice: ( + stepId: string, + pairId: string, + direction: Direction, + branchLabels: string[], + ) => void; + depth?: number; +}) { + const rendered = new Set(); + const elements: JSX.Element[] = []; + + for (let i = 0; i < steps.length; i++) { + const step = steps[i]; + if (rendered.has(step.id)) continue; + rendered.add(step.id); + + // Check if this is a send with a paired receive following it + if ( + step.type === "action" && + step.direction === "send" && + step.pairId && + i + 1 < steps.length && + steps[i + 1].pairId === step.pairId + ) { + const recvStep = steps[i + 1]; + rendered.add(recvStep.id); + elements.push( + , + ); + } else { + elements.push( + , + ); + } + } + + return <>{elements}; +} + +// ── Pair Card (!? grouped) ───────────────────────────────── + +function PairCard({ + sendStep, + recvStep, + tools, + receiveOptions, + onUpdateStep, + onRemoveStep, + onConvertToChoice, +}: { + sendStep: ProtocolStep; + recvStep: ProtocolStep; + tools: Tool[]; + receiveOptions: string[]; + onUpdateStep: ( + stepId: string, + updater: (s: ProtocolStep) => ProtocolStep, + ) => void; + onRemoveStep: (stepId: string) => void; + onConvertToChoice: ( + stepId: string, + pairId: string, + direction: Direction, + branchLabels: string[], + ) => void; +}) { + return ( +
+ {/* Send row */} +
+ + + ! + + {tools.length > 0 ? ( + + ) : ( + { + const newLabel = e.target.value; + onUpdateStep(sendStep.id, (s) => ({ ...s, label: newLabel })); + onUpdateStep(recvStep.id, (s) => ({ + ...s, + label: `${newLabel}Result`, + })); + }} + className="font-mono text-xs bg-transparent border-none outline-none flex-1 min-w-0" + placeholder="label" + /> + )} + + +
+ {/* Receive row */} +
+
+ + ? + + {receiveOptions.length > 0 ? ( + + ) : ( + + onUpdateStep(recvStep.id, (s) => ({ + ...s, + label: e.target.value, + })) + } + className="font-mono text-xs bg-transparent border-none outline-none flex-1 min-w-0" + placeholder="response label" + /> + )} + + paired +
+
+ ); +} + // ── Step Card ────────────────────────────────────────────── function StepCard({ step, - onRemove, - onUpdateLabel, - onUpdateBranchLabel, - onAddBranch, - onRemoveBranch, + tools, + recVars, + receiveOptions, + insertTarget, + onSetInsertTarget, + depth = 0, + onUpdateStep, + onRemoveStep, + onAddStepToBranch, + onConvertToChoice, }: { step: ProtocolStep; - onRemove: () => void; - onUpdateLabel: (label: string) => void; - onUpdateBranchLabel: ( + tools: Tool[]; + recVars: string[]; + receiveOptions: string[]; + insertTarget: InsertTarget | null; + onSetInsertTarget: (target: InsertTarget | null) => void; + depth?: number; + onUpdateStep: ( stepId: string, + updater: (step: ProtocolStep) => ProtocolStep, + ) => void; + onRemoveStep: (stepId: string) => void; + onAddStepToBranch: ( + choiceStepId: string, branchId: string, - label: string, + newStep: ProtocolStep, + ) => void; + onConvertToChoice: ( + stepId: string, + pairId: string, + direction: Direction, + branchLabels: string[], ) => void; - onAddBranch: () => void; - onRemoveBranch: (branchId: string) => void; }) { if (step.type === "action" && step.isRecRef) { return ( @@ -565,7 +1131,10 @@ function StepCard({ ↻ loop → {step.recVar}
-
@@ -589,19 +1158,64 @@ function StepCard({ {isSend ? "!" : "?"} - onUpdateLabel(e.target.value)} - className="font-mono text-xs bg-transparent border-none outline-none flex-1 min-w-0" - placeholder="label" - /> + {isSend && tools.length > 0 ? ( + + ) : !isSend && receiveOptions.length > 0 ? ( + + ) : ( + + onUpdateStep(step.id, (s) => ({ ...s, label: e.target.value })) + } + className="font-mono text-xs bg-transparent border-none outline-none flex-1 min-w-0" + placeholder="label" + /> + )} {step.toolName && ( ({step.toolName}) )} -
@@ -618,8 +1232,8 @@ function StepCard({ : "text-blue-600 dark:text-blue-400"; return ( -
-
+
+
{isSend ? "!" : "?"} @@ -630,45 +1244,66 @@ function StepCard({
-
+
{step.branches?.map((branch) => ( -
- - onUpdateBranchLabel(step.id, branch.id, e.target.value) + branch={branch} + choiceStep={step} + tools={tools} + recVars={recVars} + receiveOptions={receiveOptions} + insertTarget={insertTarget} + onSetInsertTarget={onSetInsertTarget} + depth={depth} + onUpdateStep={onUpdateStep} + onRemoveStep={onRemoveStep} + onAddStepToBranch={onAddStepToBranch} + onConvertToChoice={onConvertToChoice} + onRemoveBranch={() => { + if ((step.branches?.length ?? 0) > 2) { + onUpdateStep(step.id, (s) => ({ + ...s, + branches: s.branches?.filter((b) => b.id !== branch.id), + })); } - className="font-mono text-xs bg-transparent border-none outline-none flex-1 min-w-0" - placeholder="branch label" - /> - {(step.branches?.length ?? 0) > 2 && ( - - )} -
+ }} + onUpdateBranchLabel={(label) => + onUpdateStep(step.id, (s) => ({ + ...s, + branches: s.branches?.map((b) => + b.id === branch.id ? { ...b, label } : b, + ), + })) + } + /> ))}
-
+
{"}"} @@ -689,7 +1324,10 @@ function StepCard({ .
-
@@ -699,6 +1337,200 @@ function StepCard({ return null; } +// ── Branch Card (with nested steps) ──────────────────────── + +function BranchCard({ + branch, + choiceStep, + tools, + recVars, + receiveOptions, + insertTarget, + onSetInsertTarget, + depth, + onUpdateStep, + onRemoveStep, + onAddStepToBranch, + onConvertToChoice, + onRemoveBranch, + onUpdateBranchLabel, +}: { + branch: ProtocolBranch; + choiceStep: ProtocolStep; + tools: Tool[]; + recVars: string[]; + receiveOptions: string[]; + insertTarget: InsertTarget | null; + onSetInsertTarget: (target: InsertTarget | null) => void; + depth: number; + onUpdateStep: ( + stepId: string, + updater: (step: ProtocolStep) => ProtocolStep, + ) => void; + onRemoveStep: (stepId: string) => void; + onAddStepToBranch: ( + choiceStepId: string, + branchId: string, + newStep: ProtocolStep, + ) => void; + onConvertToChoice: ( + stepId: string, + pairId: string, + direction: Direction, + branchLabels: string[], + ) => void; + onRemoveBranch: () => void; + onUpdateBranchLabel: (label: string) => void; +}) { + const [expanded, setExpanded] = useState(true); + const branchTerminated = isTerminated(branch.steps); + const isActive = + insertTarget?.choiceStepId === choiceStep.id && + insertTarget?.branchId === branch.id; + + // Collect sibling branch labels to exclude from dropdown + const siblingLabels = (choiceStep.branches || []) + .filter((b) => b.id !== branch.id) + .map((b) => b.label); + + return ( +
+ {/* Branch header */} +
+ + {(() => { + const isExternal = choiceStep.direction === "receive"; + const allOptions = isExternal + ? receiveOptions + : tools.map((t) => t.name); + // Filter out labels already used by sibling branches + const options = allOptions.filter((o) => !siblingLabels.includes(o)); + if (options.length > 0 || allOptions.length > 0) { + return ( + + ); + } + return ( + onUpdateBranchLabel(e.target.value)} + className="font-mono text-xs bg-transparent border-none outline-none flex-1 min-w-0" + placeholder="branch label" + /> + ); + })()} + + {branch.steps.length > 0 && `(${branch.steps.length})`} + + {/* Target button — click to make palette add to this branch */} + + {(choiceStep.branches?.length ?? 0) > 2 && ( + + )} +
+ + {/* Branch steps */} + {expanded && ( +
+ {branch.steps.length > 0 && ( +
+ +
+ )} + + {/* Branch status */} + {branchTerminated ? ( +

end

+ ) : ( +

+ {isActive ? ( + + Use the palette to add steps here + + ) : ( + <> + Click to add steps from + palette + + )} +

+ )} +
+ )} +
+ ); +} + // ── Protocol Syntax Highlighting ─────────────────────────── function ProtocolHighlight({ protocol }: { protocol: string }) { @@ -723,7 +1555,7 @@ function ProtocolHighlight({ protocol }: { protocol: string }) { i++; } else if ( src.slice(i, i + 3) === "rec" && - (i + 3 >= src.length || !/[a-zA-Z0-9_]/.test(src[i + 3])) + (i + 3 >= src.length || !/[a-zA-Z0-9_\-]/.test(src[i + 3])) ) { tokens.push({ text: "rec", @@ -732,13 +1564,13 @@ function ProtocolHighlight({ protocol }: { protocol: string }) { i += 3; } else if ( src.slice(i, i + 3) === "end" && - (i + 3 >= src.length || !/[a-zA-Z0-9_]/.test(src[i + 3])) + (i + 3 >= src.length || !/[a-zA-Z0-9_\-]/.test(src[i + 3])) ) { tokens.push({ text: "end", cls: "text-gray-500 italic" }); i += 3; - } else if (/[a-zA-Z0-9_]/.test(src[i])) { + } else if (/[a-zA-Z0-9_\-]/.test(src[i])) { const start = i; - while (i < src.length && /[a-zA-Z0-9_]/.test(src[i])) i++; + while (i < src.length && /[a-zA-Z0-9_\-]/.test(src[i])) i++; tokens.push({ text: src.slice(start, i), cls: "text-gray-900 dark:text-gray-100", @@ -766,16 +1598,19 @@ function ProtocolHighlight({ protocol }: { protocol: string }) { // ── State Machine Preview ────────────────────────────────── function StateMachinePreview({ protocol }: { protocol: string }) { - // Simple parser → FSM for visualization - const { states, transitions } = useMemo(() => { + const { states, transitions, endStates } = useMemo(() => { try { return parseToFSM(protocol); } catch { - return { states: [], transitions: [] }; + return { + states: new Set(), + transitions: [] as FSMTransition[], + endStates: new Set(), + }; } }, [protocol]); - if (states.length === 0) { + if (states.size === 0) { return (

Add steps to see the state machine @@ -783,131 +1618,231 @@ function StateMachinePreview({ protocol }: { protocol: string }) { ); } + // Group transitions by source state for visual branching + const bySource = new Map(); + for (const t of transitions) { + const arr = bySource.get(t.from) || []; + arr.push(t); + bySource.set(t.from, arr); + } + + // Collect all states referenced in transitions + const allStates = Array.from(states).sort((a, b) => a - b); + return ( -

- {transitions.map((t, i) => ( -
- - S{t.from} - - → - - {t.dir === "send" ? "!" : t.dir === "receive" ? "?" : ""} - {t.label} - - → -
- ))} - - S{states.length > 0 ? states[states.length - 1] : 0} - +
+ {allStates.map((s) => { + const outgoing = bySource.get(s); + if (!outgoing || outgoing.length === 0) { + // Terminal state + if (endStates.has(s)) { + return ( +
+ + S{s} + + end +
+ ); + } + return null; + } + + const isBranching = outgoing.length > 1; + + return ( +
+ + S{s} + + → +
+ {outgoing.map((t, idx) => ( +
+ + {t.dir === "send" ? "!" : t.dir === "receive" ? "?" : "↻"} + {t.label} + + → + + S{t.to} + +
+ ))} +
+
+ ); + })}
); } +type FSMTransition = { from: number; to: number; dir: string; label: string }; + function parseToFSM(protocol: string): { - states: number[]; - transitions: { from: number; to: number; dir: string; label: string }[]; + states: Set; + transitions: FSMTransition[]; + endStates: Set; } { - const states: number[] = [0]; - const transitions: { - from: number; - to: number; - dir: string; - label: string; - }[] = []; - let stateId = 0; - let i = 0; + const states = new Set([0]); + const transitions: FSMTransition[] = []; + const endStates = new Set(); + let nextState = 1; + let pos = 0; const src = protocol; + const recVarStates = new Map(); const skipWS = () => { - while (i < src.length && " \t\n\r".includes(src[i])) i++; + while (pos < src.length && " \t\n\r".includes(src[pos])) pos++; }; const readIdent = () => { skipWS(); - const start = i; - while (i < src.length && /[a-zA-Z0-9_]/.test(src[i])) i++; - return src.slice(start, i); + const start = pos; + while (pos < src.length && /[a-zA-Z0-9_\-]/.test(src[pos])) pos++; + return src.slice(start, pos); }; - const parseSimple = () => { + // Returns the state number after parsing, given the current state + const parse = (currentState: number): number => { skipWS(); - if (i >= src.length) return; + if (pos >= src.length) { + endStates.add(currentState); + return currentState; + } - if (src[i] === "!" || src[i] === "?") { - const dir = src[i] === "!" ? "send" : "receive"; - i++; + if ( + src.slice(pos, pos + 3) === "end" && + (pos + 3 >= src.length || !/[a-zA-Z0-9_\-]/.test(src[pos + 3])) + ) { + pos += 3; + endStates.add(currentState); + return currentState; + } + + if ( + src.slice(pos, pos + 3) === "rec" && + (pos + 3 >= src.length || !/[a-zA-Z0-9_\-]/.test(src[pos + 3])) + ) { + pos += 3; + const varName = readIdent(); + recVarStates.set(varName, currentState); + skipWS(); + if (src[pos] === ".") pos++; + return parse(currentState); + } + + if (src[pos] === "!" || src[pos] === "?") { + const dir = src[pos] === "!" ? "send" : "receive"; + pos++; skipWS(); - if (i < src.length && src[i] === "{") { - // choice - i++; - while (i < src.length && src[i] !== "}") { + + if (pos < src.length && src[pos] === "{") { + // Choice: branches all leave from currentState + pos++; + const branchEndStates: number[] = []; + + while (pos < src.length && src[pos] !== "}") { skipWS(); + if (src[pos] === ",") { + pos++; + continue; + } + if (src[pos] === "}") break; + const label = readIdent(); - if (label) { - const from = stateId; - stateId++; - states.push(stateId); - transitions.push({ from, to: stateId, dir, label }); + if (!label) { + pos++; + continue; } + + const branchTarget = nextState++; + states.add(branchTarget); + transitions.push({ + from: currentState, + to: branchTarget, + dir, + label, + }); + skipWS(); - if (src[i] === ".") { - i++; - parseSimple(); + if (src[pos] === ".") { + pos++; + const branchEnd = parse(branchTarget); + branchEndStates.push(branchEnd); + } else { + branchEndStates.push(branchTarget); } - if (src[i] === ",") i++; + skipWS(); + if (src[pos] === ",") pos++; } - if (src[i] === "}") i++; + if (src[pos] === "}") pos++; + + // After choice, no continuation at the same level (choice terminates) + // Return the last branch end for chaining purposes, though choices shouldn't chain + return branchEndStates.length > 0 ? branchEndStates[0] : currentState; } else { + // Simple action const label = readIdent(); if (label) { - const from = stateId; - stateId++; - states.push(stateId); - transitions.push({ from, to: stateId, dir, label }); + const targetState = nextState++; + states.add(targetState); + transitions.push({ from: currentState, to: targetState, dir, label }); + + skipWS(); + if (src[pos] === ".") { + pos++; + return parse(targetState); + } + endStates.add(targetState); + return targetState; } } - } else if (src.slice(i, i + 3) === "rec") { - i += 3; - readIdent(); // var name - skipWS(); - if (src[i] === ".") i++; - parseSimple(); - return; - } else if (src.slice(i, i + 3) === "end") { - i += 3; - return; } else { - // recvar or label + // Possibly a rec var reference const ident = readIdent(); - if (ident && ident !== "end") { - // rec var reference — skip + if (ident && recVarStates.has(ident)) { + const loopTarget = recVarStates.get(ident)!; + transitions.push({ + from: currentState, + to: loopTarget, + dir: "loop", + label: ident, + }); + return currentState; } } skipWS(); - if (src[i] === ".") { - i++; - parseSimple(); + if (src[pos] === ".") { + pos++; + return parse(currentState); } + + return currentState; }; try { - parseSimple(); + parse(0); } catch { // graceful fallback } - return { states, transitions }; + return { states, transitions, endStates }; } export default ProtocolBuilderTab; diff --git a/examples/weather_agent.py b/examples/weather_agent.py new file mode 100644 index 000000000..cd8c640f9 --- /dev/null +++ b/examples/weather_agent.py @@ -0,0 +1,108 @@ +""" +Agent using MCP server-everything tools, monitored by llmsessioncontract. + +Protocol (from Protocol Builder): + !echo.?{echoResult.!get-sum.?{get-sumResult.end, get-sumError.end}, echoError.end} + +Usage: + python weather_agent.py # happy path + python weather_agent.py --echo-fail # echo fails, takes echoError branch + python weather_agent.py --sum-fail # get-sum fails, takes get-sumError branch + python weather_agent.py --violation # out-of-order call +""" + +import json +import sys + +from llmcontract import Monitor + +# ── Protocol (as generated by Protocol Builder) ───────────── +PROTOCOL = "!echo.?{echoResult.!get-sum.?{get-sumResult.end, get-sumError.end}, echoError.end}" + +# ── Tool implementations (mimic MCP server-everything) ────── + +def echo(message: str) -> dict: + return {"echo": message} + +def get_sum(a: float, b: float) -> dict: + return {"sum": a + b} + +# ── Monitored agent ───────────────────────────────────────── + +def run_agent(echo_fails=False, sum_fails=False): + print(f"\nProtocol: {PROTOCOL}") + print(f"{'='*60}\n") + + monitor = Monitor(PROTOCOL) + + # Step 1: !echo + print(f"[S{monitor.current_state}] Calling echo tool...") + r = monitor.send("echo") + print(f" !echo -> {r}") + + if echo_fails: + print(f"\n[S{monitor.current_state}] Echo failed!") + r = monitor.receive("echoError") + print(f" ?echoError -> {r}") + print(f" Terminal: {monitor.is_terminal}") + return + + result = echo("Hello from the monitored agent!") + print(f" Result: {result}") + + # Step 2: ?echoResult + r = monitor.receive("echoResult") + print(f" ?echoResult -> {r}") + + # Step 3: !get-sum + print(f"\n[S{monitor.current_state}] Calling get-sum tool...") + r = monitor.send("get-sum") + print(f" !get-sum -> {r}") + + if sum_fails: + print(f"\n[S{monitor.current_state}] get-sum failed!") + r = monitor.receive("get-sumError") + print(f" ?get-sumError -> {r}") + print(f" Terminal: {monitor.is_terminal}") + return + + result = get_sum(3, 5) + print(f" Result: {result}") + + # Step 4: ?get-sumResult + r = monitor.receive("get-sumResult") + print(f" ?get-sumResult -> {r}") + + print(f"\n[S{monitor.current_state}] Done! Terminal: {monitor.is_terminal}") + + +def demo_violation(): + print(f"\nProtocol: {PROTOCOL}") + print(f"{'='*60}") + print("Attempting to call get-sum before echo completes...\n") + + monitor = Monitor(PROTOCOL) + + r = monitor.send("echo") + print(f"[S{monitor.current_state}] !echo -> {r}") + + # VIOLATION: skip ?echoResult and try !get-sum directly + r = monitor.send("get-sum") + print(f"[S{monitor.current_state}] !get-sum -> {r}") + print(f" Halted: {monitor.is_halted}") + + r = monitor.receive("echoResult") + print(f" ?echoResult -> {r}") + + +if __name__ == "__main__": + args = sys.argv[1:] + + if "--violation" in args: + demo_violation() + elif "--echo-fail" in args: + run_agent(echo_fails=True) + elif "--sum-fail" in args: + run_agent(sum_fails=True) + else: + run_agent()