diff --git a/client/src/App.tsx b/client/src/App.tsx index 59d15ba06..497e29c3c 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; @@ -480,6 +483,7 @@ const App = () => { "roots", "auth", "metadata", + "protocol-builder", ]; const isValidTab = validTabs.includes(hash); @@ -814,6 +818,7 @@ const App = () => { "roots", "auth", "metadata", + "protocol-builder", ]; if (validTabs.includes(originatingTab)) { @@ -1437,6 +1442,10 @@ const App = () => { Metadata + + + Protocol Builder +
@@ -1673,6 +1682,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..8f242d334 --- /dev/null +++ b/client/src/components/ProtocolBuilderTab.tsx @@ -0,0 +1,1848 @@ +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, + Target, + GitBranch, +} 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; + pairId?: string; // links ! and ? steps together +} + +interface ProtocolBranch { + id: string; + label: string; + steps: ProtocolStep[]; +} + +// Target for where palette clicks should add steps +interface InsertTarget { + choiceStepId: string; + branchId: string; +} + +// ── 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 llmsessioncontract 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,'} + }, +)`; +} + +// ── 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 = ({ + tools, + listTools, +}: { + tools: Tool[]; + listTools: () => void; +}) => { + 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"; + return stepsToProtocol(steps); + }, [steps]); + + const pythonSnippet = useMemo( + () => generatePythonSnippet(protocol, tools), + [protocol, tools], + ); + + 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 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: [] }, + ], + }; + addStepToTarget(step); + }, + [addStepToTarget], + ); + + const addRecursion = useCallback(() => { + 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 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 handleConvertToChoice = useCallback( + ( + stepId: string, + pairId: string, + direction: Direction, + branchLabels: string[], + ) => { + setSteps((prev) => + convertPairToChoice(prev, stepId, pairId, direction, branchLabels), + ); + }, + [], + ); + + const handleAddStepToBranch = useCallback( + (choiceStepId: string, branchId: string, newStep: ProtocolStep) => { + setSteps((prev) => + addStepToBranch(prev, choiceStepId, branchId, newStep), + ); + }, + [], + ); + + const copyToClipboard = useCallback( + (text: string, type: "dsl" | "python") => { + navigator.clipboard.writeText(text); + setCopied(type); + setTimeout(() => setCopied(null), 2000); + }, + [], + ); + + const clearAll = useCallback(() => { + setSteps([]); + setInsertTarget(null); + }, []); + + 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 ( + +
+
+ {/* ── Left: Tool Palette ── */} +
+
+

+ MCP Tools +

+ +
+ + {/* Insert target indicator */} + {insertTarget && ( +
+ + + Adding to: {targetLabel} + + +
+ )} + + {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 +

+
+ ) : ( +
+ + + {/* Terminal indicator */} +
+ + end + +
+
+ )} +
+ + {/* ── Right: Output ── */} +
+

+ Output +

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

+ Session Type DSL +

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

+ State Machine Preview +

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

+ Python Integration +

+ +
+
+                {pythonSnippet}
+              
+
+ + {/* Download button */} + +
+
+
+
+ ); +}; + +// ── 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, + tools, + recVars, + receiveOptions, + insertTarget, + onSetInsertTarget, + depth = 0, + onUpdateStep, + onRemoveStep, + onAddStepToBranch, + onConvertToChoice, +}: { + step: 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; +}) { + 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 ? "!" : "?"} + + {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}) + + )} + +
+ ); + } + + 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) => ( + { + if ((step.branches?.length ?? 0) > 2) { + onUpdateStep(step.id, (s) => ({ + ...s, + branches: s.branches?.filter((b) => b.id !== branch.id), + })); + } + }} + onUpdateBranchLabel={(label) => + onUpdateStep(step.id, (s) => ({ + ...s, + branches: s.branches?.map((b) => + b.id === branch.id ? { ...b, label } : b, + ), + })) + } + /> + ))} +
+
+ + {"}"} + +
+
+ ); + } + + if (step.type === "recursion") { + return ( +
+ + + rec + + + {step.recVar} + + . +
+ +
+ ); + } + + 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 }) { + 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 }) { + const { states, transitions, endStates } = useMemo(() => { + try { + return parseToFSM(protocol); + } catch { + return { + states: new Set(), + transitions: [] as FSMTransition[], + endStates: new Set(), + }; + } + }, [protocol]); + + if (states.size === 0) { + return ( +

+ Add steps to see the state machine +

+ ); + } + + // 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 ( +
+ {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: Set; + transitions: FSMTransition[]; + endStates: Set; +} { + 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 (pos < src.length && " \t\n\r".includes(src[pos])) pos++; + }; + + const readIdent = () => { + skipWS(); + const start = pos; + while (pos < src.length && /[a-zA-Z0-9_\-]/.test(src[pos])) pos++; + return src.slice(start, pos); + }; + + // Returns the state number after parsing, given the current state + const parse = (currentState: number): number => { + skipWS(); + if (pos >= src.length) { + endStates.add(currentState); + return currentState; + } + + 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 (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) { + pos++; + continue; + } + + const branchTarget = nextState++; + states.add(branchTarget); + transitions.push({ + from: currentState, + to: branchTarget, + dir, + label, + }); + + skipWS(); + if (src[pos] === ".") { + pos++; + const branchEnd = parse(branchTarget); + branchEndStates.push(branchEnd); + } else { + branchEndStates.push(branchTarget); + } + skipWS(); + if (src[pos] === ",") pos++; + } + 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 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 { + // Possibly a rec var reference + const ident = readIdent(); + 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[pos] === ".") { + pos++; + return parse(currentState); + } + + return currentState; + }; + + try { + parse(0); + } catch { + // graceful fallback + } + + 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()