fix(examples): improve updateModelContext with structured YAML frontmatter (#271)

* map-server: fix async/await in ontoolresult, remove unused getViewStorageKey

* fix(map-server): improve updateModelContext and fix initialization

- Use structured YAML frontmatter format for model context (like pdf-server)
- Add tool name, altitude, and cleaner data format
- Fix initialization deadlock: app.connect() was inside ontoolresult handler
  which could never fire until connected
- Fix view state persistence: widgetUUID now properly extracted from
  ontoolresult and used to restore persisted camera position
- Server now includes widgetUUID in tool result _meta for widget identification

* fix(pdf-server): use widgetUUID for page persistence

- Add widgetUUID to tool result _meta for widget identification
- Use widgetUUID instead of generated key (pdfUrl+toolId) for localStorage
- Enables reliable page restoration when revisiting PDF widgets

* fix(transcript-server): use structured YAML frontmatter for model context

- Add tool name, listening status, and unsent entry count in frontmatter
- Consistent format with pdf-server and map-server

* update names

* drop bogus altitude

* Update mcp-app.ts

Author: ochafik

examples/map-server/server.ts

@@ -21,6 +21,7 @@ import {
   RESOURCE_URI_META_KEY,
 } from "@modelcontextprotocol/ext-apps/server";
 import { startServer } from "./server-utils.js";
+import { randomUUID } from "crypto";
 
 const DIST_DIR = path.join(import.meta.dirname, "dist");
 const RESOURCE_URI = "ui://cesium-map/mcp-app.html";
@@ -182,6 +183,9 @@ export function createServer(): McpServer {
           text: `Displaying globe at: W:${west.toFixed(4)}, S:${south.toFixed(4)}, E:${east.toFixed(4)}, N:${north.toFixed(4)}${label ? ` (${label})` : ""}`,
         },
       ],
+      _meta: {
+        widgetUUID: randomUUID(),
+      },
     }),
   );
 

examples/map-server/src/mcp-app.ts

@@ -71,6 +71,8 @@ let persistViewTimer: ReturnType<typeof setTimeout> | null = null;
 // Track whether tool input has been received (to know if we should restore persisted state)
 let hasReceivedToolInput = false;
 
+let widgetUUID: string | undefined = undefined;
+
 /**
  * Persisted camera state for localStorage
  */
@@ -83,18 +85,6 @@ interface PersistedCameraState {
   roll: number; // radians
 }
 
-/**
- * Get localStorage key for persisting view state
- * Uses toolInfo.id (tool invocation ID) - localStorage is scoped per conversation per server,
- * so each tool call remembers its own view state within the conversation.
- */
-function getViewStorageKey(): string | null {
-  const context = app.getHostContext();
-  const toolId = context?.toolInfo?.id;
-  if (!toolId) return null;
-  return `cesium-view:${toolId}`;
-}
-
 /**
  * Get current camera state for persistence
  */
@@ -132,8 +122,7 @@ function schedulePersistViewState(cesiumViewer: any): void {
  * Persist current view state to localStorage
  */
 function persistViewState(cesiumViewer: any): void {
-  const key = getViewStorageKey();
-  if (!key) {
+  if (!widgetUUID) {
     log.info("No storage key available, skipping view persistence");
     return;
   }
@@ -142,13 +131,9 @@ function persistViewState(cesiumViewer: any): void {
   if (!state) return;
 
   try {
-    localStorage.setItem(key, JSON.stringify(state));
-    log.info(
-      "Persisted view state:",
-      key,
-      state.latitude.toFixed(2),
-      state.longitude.toFixed(2),
-    );
+    const value = JSON.stringify(state);
+    localStorage.setItem(widgetUUID, value);
+    log.info("Persisted view state:", widgetUUID, value);
   } catch (e) {
     log.warn("Failed to persist view state:", e);
   }
@@ -158,12 +143,14 @@ function persistViewState(cesiumViewer: any): void {
  * Load persisted view state from localStorage
  */
 function loadPersistedViewState(): PersistedCameraState | null {
-  const key = getViewStorageKey();
-  if (!key) return null;
+  if (!widgetUUID) return null;
 
   try {
-    const stored = localStorage.getItem(key);
-    if (!stored) return null;
+    const stored = localStorage.getItem(widgetUUID);
+    if (!stored) {
+      console.info("No persisted view state found");
+      return null;
+    }
 
     const state = JSON.parse(stored) as PersistedCameraState;
     // Basic validation
@@ -175,6 +162,7 @@ function loadPersistedViewState(): PersistedCameraState | null {
       log.warn("Invalid persisted view state, ignoring");
       return null;
     }
+    log.info("Loaded persisted view state:", state);
     return state;
   } catch (e) {
     log.warn("Failed to load persisted view state:", e);
@@ -402,8 +390,10 @@ async function getVisiblePlaces(extent: BoundingBox): Promise<string[]> {
 }
 
 /**
- * Debounced location update using multi-point reverse geocoding
- * Samples multiple points in the visible extent to discover places
+ * Debounced location update using multi-point reverse geocoding.
+ * Samples multiple points in the visible extent to discover places.
+ *
+ * Updates model context with structured YAML frontmatter (similar to pdf-server).
  */
 function scheduleLocationUpdate(cesiumViewer: any): void {
   if (reverseGeocodeTimer) {
@@ -420,34 +410,35 @@ function scheduleLocationUpdate(cesiumViewer: any): void {
     }
 
     const { widthKm, heightKm } = getScaleDimensions(extent);
-    const extentInfo =
-      `Extent: [${extent.west.toFixed(4)}, ${extent.south.toFixed(4)}, ` +
-      `${extent.east.toFixed(4)}, ${extent.north.toFixed(4)}] ` +
-      `(${widthKm.toFixed(1)}km × ${heightKm.toFixed(1)}km)`;
-    log.info(extentInfo);
+
+    log.info(`Extent: ${widthKm.toFixed(1)}km × ${heightKm.toFixed(1)}km`);
 
     // Get places visible in the extent (samples multiple points for large areas)
     const places = await getVisiblePlaces(extent);
-    const placesText =
-      places.length > 0 ? `Visible places: ${places.join(", ")}` : "";
-
-    if (places.length > 0 || center) {
-      const centerText = center
-        ? `Center: ${center.lat.toFixed(4)}, ${center.lon.toFixed(4)}`
-        : "";
-
-      const contextText = [placesText, centerText, extentInfo]
-        .filter(Boolean)
-        .join("\n");
 
-      log.info("Updating model context:", contextText);
-
-      // Update the model's context with the current map location.
-      // If the host doesn't support this, the request will silently fail.
-      app.updateModelContext({
-        content: [{ type: "text", text: contextText }],
-      });
-    }
+    // Build structured markdown with YAML frontmatter (like pdf-server)
+    // Note: tool name isn't in the notification protocol, so we hardcode it
+    const frontmatter = [
+      "---",
+      `tool: show-map`,
+      center
+        ? `center: [${center.lat.toFixed(4)}, ${center.lon.toFixed(4)}]`
+        : null,
+      `extent: [${extent.west.toFixed(4)}, ${extent.south.toFixed(4)}, ${extent.east.toFixed(4)}, ${extent.north.toFixed(4)}]`,
+      `extent-size: ${widthKm.toFixed(1)}km × ${heightKm.toFixed(1)}km`,
+      places.length > 0 ? `visible-places: [${places.join(", ")}]` : null,
+      "---",
+    ]
+      .filter(Boolean)
+      .join("\n");
+
+    log.info("Updating model context:", frontmatter);
+
+    // Update the model's context with the current map location.
+    // If the host doesn't support this, the request will silently fail.
+    app.updateModelContext({
+      content: [{ type: "text", text: frontmatter }],
+    });
   }, 1500);
 }
 
@@ -947,40 +938,36 @@ app.ontoolinput = async (params) => {
 //   },
 // );
 
+// Handle tool result - extract widgetUUID and restore persisted view if available
+app.ontoolresult = async (result) => {
+  widgetUUID = result._meta?.widgetUUID
+    ? String(result._meta.widgetUUID)
+    : undefined;
+  log.info("Tool result received, widgetUUID:", widgetUUID);
+
+  // Now that we have widgetUUID, try to restore persisted view
+  // This overrides the tool input position if a saved state exists
+  if (viewer && widgetUUID) {
+    const restored = restorePersistedView(viewer);
+    if (restored) {
+      log.info("Restored persisted view from tool result handler");
+      await waitForTilesLoaded(viewer);
+      hideLoading();
+    }
+  }
+};
+
 // Initialize Cesium and connect to host
-async function init() {
+async function initialize() {
   try {
     log.info("Loading CesiumJS from CDN...");
     await loadCesium();
     log.info("CesiumJS loaded successfully");
 
     viewer = await initCesium();
-    // Don't hide loading here - we wait for tool input to position camera
-    // and for tiles to load before hiding the loading indicator
-    log.info("CesiumJS initialized, waiting for tool input...");
+    log.info("CesiumJS initialized");
 
-    // Fallback: if no tool input received within 2 seconds, try restoring
-    // persisted view or show default view
-    setTimeout(async () => {
-      const loadingEl = document.getElementById("loading");
-      if (
-        loadingEl &&
-        loadingEl.style.display !== "none" &&
-        !hasReceivedToolInput
-      ) {
-        // No explicit tool input - try to restore persisted view
-        const restored = restorePersistedView(viewer!);
-        if (restored) {
-          log.info("Restored persisted view, waiting for tiles...");
-        } else {
-          log.info("No persisted view, using default view...");
-        }
-        await waitForTilesLoaded(viewer!);
-        hideLoading();
-      }
-    }, 2000);
-
-    // Connect to host (auto-creates PostMessageTransport)
+    // Connect to host (must happen before we can receive notifications)
     await app.connect();
     log.info("Connected to host");
 
@@ -1009,6 +996,26 @@ async function init() {
 
     // Set up keyboard shortcuts for fullscreen (Escape to exit, Ctrl/Cmd+Enter to toggle)
     document.addEventListener("keydown", handleFullscreenKeyboard);
+
+    // Wait a bit for tool input, then try restoring persisted view or show default
+    setTimeout(async () => {
+      const loadingEl = document.getElementById("loading");
+      if (
+        loadingEl &&
+        loadingEl.style.display !== "none" &&
+        !hasReceivedToolInput
+      ) {
+        // No explicit tool input - try to restore persisted view
+        const restored = restorePersistedView(viewer!);
+        if (restored) {
+          log.info("Restored persisted view, waiting for tiles...");
+        } else {
+          log.info("No persisted view, using default view...");
+        }
+        await waitForTilesLoaded(viewer!);
+        hideLoading();
+      }
+    }, 500);
   } catch (error) {
     log.error("Failed to initialize:", error);
     const loadingEl = document.getElementById("loading");
@@ -1019,4 +1026,5 @@ async function init() {
   }
 }
 
-init();
+// Start initialization
+initialize();

examples/pdf-server/server.ts

@@ -21,6 +21,7 @@ import type {
 import fs from "node:fs/promises";
 import path from "node:path";
 import { z } from "zod";
+import { randomUUID } from "crypto";
 
 import {
   buildPdfIndex,
@@ -175,6 +176,9 @@ The viewer supports zoom, navigation, text selection, and fullscreen mode.`,
           },
         ],
         structuredContent: result,
+        _meta: {
+          widgetUUID: randomUUID(),
+        },
       };
     },
   );

examples/pdf-server/src/mcp-app.ts

@@ -35,6 +35,7 @@ let totalPages = 0;
 let scale = 1.0;
 let pdfUrl = "";
 let pdfTitle: string | undefined;
+let widgetUUID: string | undefined;
 let currentRenderTask: { cancel: () => void } | null = null;
 
 // DOM Elements
@@ -402,20 +403,11 @@ async function renderPage() {
   }
 }
 
-// Page persistence
-function getStorageKey(): string | null {
-  if (!pdfUrl) return null;
-  const ctx = app.getHostContext();
-  const toolId = ctx?.toolInfo?.id ?? pdfUrl;
-  return `pdf:${pdfUrl}:${toolId}`;
-}
-
 function saveCurrentPage() {
-  const key = getStorageKey();
-  log.info("saveCurrentPage: key=", key, "page=", currentPage);
-  if (key) {
+  log.info("saveCurrentPage: key=", widgetUUID, "page=", currentPage);
+  if (widgetUUID) {
     try {
-      localStorage.setItem(key, String(currentPage));
+      localStorage.setItem(widgetUUID, String(currentPage));
       log.info("saveCurrentPage: saved successfully");
     } catch (err) {
       log.error("saveCurrentPage: error", err);
@@ -424,11 +416,10 @@ function saveCurrentPage() {
 }
 
 function loadSavedPage(): number | null {
-  const key = getStorageKey();
-  log.info("loadSavedPage: key=", key);
-  if (!key) return null;
+  log.info("loadSavedPage: key=", widgetUUID);
+  if (!widgetUUID) return null;
   try {
-    const saved = localStorage.getItem(key);
+    const saved = localStorage.getItem(widgetUUID);
     log.info("loadSavedPage: saved value=", saved);
     if (saved) {
       const page = parseInt(saved, 10);
@@ -715,6 +706,9 @@ app.ontoolresult = async (result) => {
   pdfUrl = parsed.url;
   pdfTitle = parsed.title;
   totalPages = parsed.pageCount;
+  widgetUUID = result._meta?.widgetUUID
+    ? String(result._meta.widgetUUID)
+    : undefined;
 
   // Restore saved page or use initial page
   const savedPage = loadSavedPage();

examples/transcript-server/src/mcp-app.ts

@@ -329,18 +329,31 @@ function updateSendButton() {
   sendBtn.disabled = unsentEntries.length === 0;
 }
 
+/**
+ * Update model context with structured YAML frontmatter (like pdf-server, map-server).
+ */
 function updateModelContext() {
   const caps = app.getHostCapabilities();
   if (!caps?.updateModelContext) return;
 
   const text = getUnsentText();
+  const unsentCount = getUnsentEntries().length;
   log.info("Updating model context:", text || "(empty)");
 
+  // Build structured markdown with YAML frontmatter
+  const frontmatter = [
+    "---",
+    "tool: transcribe",
+    `status: ${isListening ? "listening" : "paused"}`,
+    `unsent-entries: ${unsentCount}`,
+    "---",
+  ].join("\n");
+
+  const markdown = text ? `${frontmatter}\n\n${text}` : frontmatter;
+
   app
     .updateModelContext({
-      content: text
-        ? [{ type: "text", text: `[Live transcript]: ${text}` }]
-        : [],
+      content: [{ type: "text", text: markdown }],
     })
     .catch((e: unknown) => {
       log.warn("Failed to update model context:", e);

tests/e2e/generate-grid-screenshots.spec.ts

@@ -52,7 +52,7 @@ const SERVERS = [
     name: "Customer Segmentation Server",
     dir: "customer-segmentation-server",
   },
-  { key: "map-server", name: "CesiumJS Map Server", dir: "map-server" },
+  { key: "map-server", name: "Map Server", dir: "map-server" },
   { key: "pdf-server", name: "PDF Server", dir: "pdf-server" },
   {
     key: "scenario-modeler",

tests/e2e/servers.spec.ts

@@ -47,7 +47,7 @@ const SERVERS = [
   { key: "budget-allocator", name: "Budget Allocator Server" },
   { key: "cohort-heatmap", name: "Cohort Heatmap Server" },
   { key: "customer-segmentation", name: "Customer Segmentation Server" },
-  { key: "map-server", name: "CesiumJS Map Server" },
+  { key: "map-server", name: "Map Server" },
   { key: "pdf-server", name: "PDF Server" },
   { key: "scenario-modeler", name: "SaaS Scenario Modeler" },
   { key: "shadertoy", name: "ShaderToy Server" },