feat(apps): add support for MCP apps to sample (#7039)

Co-authored-by: Andrew Harvard <aharvard@block.xyz>

crates/goose-server/src/routes/mod.rs

@@ -9,6 +9,7 @@ pub mod prompts;
 pub mod recipe;
 pub mod recipe_utils;
 pub mod reply;
+pub mod sampling;
 pub mod schedule;
 pub mod session;
 pub mod setup;
@@ -39,4 +40,5 @@ pub fn configure(state: Arc<crate::state::AppState>, secret_key: String) -> Rout
         .merge(tunnel::routes(state.clone()))
         .merge(mcp_ui_proxy::routes(secret_key.clone()))
         .merge(mcp_app_proxy::routes(secret_key))
+        .merge(sampling::routes(state))
 }

crates/goose-server/src/routes/sampling.rs

@@ -0,0 +1,87 @@
+use axum::{
+    extract::{Path, State},
+    http::StatusCode,
+    routing::post,
+    Json, Router,
+};
+use goose::conversation::message::Message;
+use rmcp::model::{
+    CreateMessageRequestParams, CreateMessageResult, Role, SamplingContent, SamplingMessage,
+    SamplingMessageContent,
+};
+use std::sync::Arc;
+
+use crate::state::AppState;
+
+pub fn routes(state: Arc<AppState>) -> Router {
+    Router::new()
+        .route(
+            "/sessions/{session_id}/sampling/message",
+            post(create_message),
+        )
+        .with_state(state)
+}
+
+async fn create_message(
+    State(state): State<Arc<AppState>>,
+    Path(session_id): Path<String>,
+    Json(request): Json<CreateMessageRequestParams>,
+) -> Result<Json<CreateMessageResult>, StatusCode> {
+    let agent = state.get_agent_for_route(session_id.clone()).await?;
+
+    let provider = agent.provider().await.map_err(|e| {
+        tracing::error!("Failed to get provider: {}", e);
+        StatusCode::INTERNAL_SERVER_ERROR
+    })?;
+
+    let messages: Vec<Message> = request
+        .messages
+        .iter()
+        .map(|msg| {
+            let base = match msg.role {
+                Role::User => Message::user(),
+                Role::Assistant => Message::assistant(),
+            };
+            content_to_message(base, &msg.content)
+        })
+        .collect();
+
+    let system = request
+        .system_prompt
+        .as_deref()
+        .unwrap_or("You are a helpful AI assistant.");
+
+    let model_config = provider.get_model_config();
+    let (response, usage) = provider
+        .complete(&model_config, &session_id, system, &messages, &[])
+        .await
+        .map_err(|e| {
+            tracing::error!("Sampling completion failed: {}", e);
+            StatusCode::INTERNAL_SERVER_ERROR
+        })?;
+
+    let text = response.as_concat_text();
+
+    Ok(Json(CreateMessageResult {
+        model: usage.model,
+        stop_reason: Some(CreateMessageResult::STOP_REASON_END_TURN.to_string()),
+        message: SamplingMessage::new(Role::Assistant, SamplingMessageContent::text(&text)),
+    }))
+}
+
+fn content_to_message(base: Message, content: &SamplingContent<SamplingMessageContent>) -> Message {
+    let items = match content {
+        SamplingContent::Single(item) => vec![item],
+        SamplingContent::Multiple(items) => items.iter().collect(),
+    };
+
+    let mut msg = base;
+    for item in items {
+        msg = match item {
+            SamplingMessageContent::Text(text) => msg.with_text(&text.text),
+            SamplingMessageContent::Image(image) => msg.with_image(&image.data, &image.mime_type),
+            _ => msg,
+        };
+    }
+    msg
+}

crates/goose/src/goose_apps/cache.rs

@@ -7,6 +7,7 @@ use tracing::warn;
 use super::app::GooseApp;
 
 static CLOCK_HTML: &str = include_str!("../goose_apps/clock.html");
+static CHAT_HTML: &str = include_str!("../goose_apps/chat.html");
 const APPS_EXTENSION_NAME: &str = "apps";
 
 pub struct McpAppCache {
@@ -23,10 +24,12 @@ impl McpAppCache {
     }
 
     fn ensure_default_apps(&self) {
-        if self.get_app(APPS_EXTENSION_NAME, "apps://clock").is_none() {
-            if let Ok(mut clock_app) = GooseApp::from_html(CLOCK_HTML) {
-                clock_app.mcp_servers = vec![APPS_EXTENSION_NAME.to_string()];
-                let _ = self.store_app(&clock_app);
+        for (uri, html) in [("apps://clock", CLOCK_HTML), ("apps://chat", CHAT_HTML)] {
+            if self.get_app(APPS_EXTENSION_NAME, uri).is_none() {
+                if let Ok(mut app) = GooseApp::from_html(html) {
+                    app.mcp_servers = vec![APPS_EXTENSION_NAME.to_string()];
+                    let _ = self.store_app(&app);
+                }
             }
         }
     }

crates/goose/src/goose_apps/chat.html

@@ -0,0 +1,184 @@
+<!DOCTYPE html>
+<html>
+<head>
+  <meta charset="UTF-8">
+  <title>Chat</title>
+  <script type="application/ld+json">
+    {
+      "@context": "https://goose.ai/schema",
+      "@type": "GooseApp",
+      "name": "chat",
+      "description": "Simple Chat UI",
+      "width": 400,
+      "height": 500,
+      "resizable": true
+    }
+  </script>
+  <style>
+    * { box-sizing: border-box; margin: 0; padding: 0; }
+    html, body { height: 100%; font-family: -apple-system, BlinkMacSystemFont, sans-serif; }
+    body { display: flex; flex-direction: column; background: #fff; }
+    
+    .messages {
+      flex: 1;
+      overflow-y: auto;
+      padding: 16px;
+      display: flex;
+      flex-direction: column;
+      gap: 12px;
+    }
+    
+    .message {
+      max-width: 80%;
+      padding: 10px 14px;
+      border-radius: 16px;
+      line-height: 1.4;
+      font-size: 14px;
+      word-wrap: break-word;
+    }
+    
+    .message.user {
+      align-self: flex-end;
+      background: #000;
+      color: #fff;
+    }
+    
+    .message.assistant {
+      align-self: flex-start;
+      background: #f0f0f0;
+      color: #000;
+    }
+    
+    .message.loading {
+      font-style: italic;
+      color: #666;
+    }
+    
+    .input-area {
+      display: flex;
+      gap: 8px;
+      padding: 12px;
+      border-top: 1px solid #e0e0e0;
+      background: #fafafa;
+    }
+    
+    #messageInput {
+      flex: 1;
+      padding: 10px 14px;
+      border: 1px solid #ddd;
+      border-radius: 20px;
+      font-size: 14px;
+      outline: none;
+    }
+    
+    #messageInput:focus { border-color: #999; }
+    
+    #sendBtn {
+      padding: 10px 20px;
+      background: #000;
+      color: #fff;
+      border: none;
+      border-radius: 20px;
+      font-size: 14px;
+      cursor: pointer;
+    }
+    
+    #sendBtn:disabled {
+      background: #ccc;
+      cursor: not-allowed;
+    }
+  </style>
+</head>
+<body>
+  <div class="messages" id="messages"></div>
+  <div class="input-area">
+    <input type="text" id="messageInput" placeholder="Type a message..." />
+    <button id="sendBtn">Send</button>
+  </div>
+
+  <script>
+    const messagesEl = document.getElementById('messages');
+    const inputEl = document.getElementById('messageInput');
+    const sendBtn = document.getElementById('sendBtn');
+    const conversationHistory = [];
+    const pendingRequests = new Map();
+    let requestId = 0;
+
+    function addMessage(role, text, isLoading = false) {
+      const div = document.createElement('div');
+      div.className = `message ${role}${isLoading ? ' loading' : ''}`;
+      div.textContent = text;
+      messagesEl.appendChild(div);
+      messagesEl.scrollTop = messagesEl.scrollHeight;
+      return div;
+    }
+
+    function request(method, params) {
+      return new Promise((resolve, reject) => {
+        const id = ++requestId;
+        pendingRequests.set(id, { resolve, reject });
+        window.parent.postMessage({ jsonrpc: '2.0', id, method, params }, '*');
+      });
+    }
+
+    window.addEventListener('message', (event) => {
+      const data = event.data;
+      if (!data || typeof data !== 'object') return;
+      
+      if ('id' in data && pendingRequests.has(data.id)) {
+        const { resolve, reject } = pendingRequests.get(data.id);
+        pendingRequests.delete(data.id);
+        if (data.error) {
+          reject(new Error(data.error.message || 'Unknown error'));
+        } else {
+          resolve(data.result);
+        }
+      }
+    });
+
+    request('ui/initialize', {}).then(() => {
+      window.parent.postMessage({ jsonrpc: '2.0', method: 'ui/notifications/initialized', params: {} }, '*');
+    });
+
+    async function sendMessage() {
+      const text = inputEl.value.trim();
+      if (!text) return;
+
+      inputEl.value = '';
+      sendBtn.disabled = true;
+
+      addMessage('user', text);
+      conversationHistory.push({ role: 'user', content: { type: 'text', text } });
+
+      const loadingEl = addMessage('assistant', 'Thinking...', true);
+
+      try {
+        const response = await request('sampling/createMessage', {
+          messages: conversationHistory,
+          systemPrompt: 'You are a helpful assistant. Keep responses concise.',
+          maxTokens: 1000
+        });
+
+        const responseText = response.content.text;
+        conversationHistory.push({ role: 'assistant', content: { type: 'text', text: responseText } });
+        
+        loadingEl.textContent = responseText;
+        loadingEl.classList.remove('loading');
+      } catch (err) {
+        loadingEl.textContent = 'Error: ' + err.message;
+        loadingEl.classList.remove('loading');
+      }
+
+      sendBtn.disabled = false;
+      inputEl.focus();
+    }
+
+    sendBtn.addEventListener('click', sendMessage);
+    inputEl.addEventListener('keypress', (e) => {
+      if (e.key === 'Enter') sendMessage();
+    });
+
+    inputEl.focus();
+  </script>
+</body>
+</html>

ui/desktop/src/components/McpApps/McpAppRenderer.tsx

@@ -40,6 +40,8 @@ import {
   McpAppToolInputPartial,
   McpAppToolResult,
   DimensionLayout,
+  SamplingCreateMessageParams,
+  SamplingCreateMessageResponse,
 } from './types';
 
 const DEFAULT_IFRAME_HEIGHT = 200;
@@ -185,6 +187,7 @@ function appReducer(state: AppState, action: AppAction): AppState {
 
   switch (action.type) {
     case 'FETCH_RESOURCE':
+      if (state.status === 'ready') return state;
       return { status: 'loading_resource', html, meta };
 
     case 'RESOURCE_LOADED':
@@ -237,81 +240,157 @@ export default function McpAppRenderer({
 
   const { resolvedTheme } = useTheme();
 
-  const initialState: AppState = cachedHtml
-    ? { status: 'loading_sandbox', html: cachedHtml, meta: DEFAULT_META }
-    : { status: 'idle' };
+  // Survive StrictMode remounts — replay cached results instead of re-fetching,
+  // which prevents the iframe from being torn down and recreated (visible flicker).
+  // Declared before useReducer so the lazy initializer can read them.
+  const fetchedDataRef = useRef<{ html: string; meta: ResourceMeta } | null>(null);
+  const sandboxUrlRef = useRef<{ url: string; csp: McpUiResourceCsp | null } | null>(null);
 
-  const [state, dispatch] = useReducer(appReducer, initialState);
+  const [state, dispatch] = useReducer(appReducer, undefined, (): AppState => {
+    // On StrictMode remount, skip straight to ready if we have all cached data.
+    if (fetchedDataRef.current && sandboxUrlRef.current) {
+      return {
+        status: 'ready',
+        html: fetchedDataRef.current.html,
+        meta: fetchedDataRef.current.meta,
+        sandboxUrl: new URL(sandboxUrlRef.current.url),
+        sandboxCsp: sandboxUrlRef.current.csp,
+      };
+    }
+    if (cachedHtml) {
+      return { status: 'loading_sandbox', html: cachedHtml, meta: DEFAULT_META };
+    }
+    return { status: 'idle' };
+  });
   const [iframeHeight, setIframeHeight] = useState(DEFAULT_IFRAME_HEIGHT);
 
   const containerRef = useRef<HTMLDivElement>(null);
   const [containerWidth, setContainerWidth] = useState<number>(0);
   const [containerHeight, setContainerHeight] = useState<number>(0);
+  const [apiHost, setApiHost] = useState<string | null>(null);
+  const [secretKey, setSecretKey] = useState<string | null>(null);
+
+  useEffect(() => {
+    window.electron.getGoosedHostPort().then(setApiHost);
+    window.electron.getSecretKey().then(setSecretKey);
+  }, []);
 
   // Fetch the resource from the extension to get HTML and metadata (CSP, permissions, etc.).
   // If cachedHtml is provided we show it immediately; the fetch updates metadata and
   // replaces HTML only if the server returns different content.
+  //
+  // Retries with exponential backoff when the fetch fails (e.g. the extension hasn't
+  // finished loading yet, causing a transient 500). Cached HTML skips retries since
+  // the app can render immediately with the cached version.
   useEffect(() => {
     if (!sessionId) return;
 
+    // On StrictMode remount, replay the cached result instead of re-fetching.
+    if (fetchedDataRef.current) {
+      const { html: cachedResult, meta: cachedMeta } = fetchedDataRef.current;
+      dispatch({ type: 'RESOURCE_LOADED', html: cachedResult, meta: cachedMeta });
+      return;
+    }
+
+    const MAX_RETRIES = 5;
+    const BASE_DELAY_MS = 500;
+    let cancelled = false;
+
     const fetchResourceData = async () => {
       dispatch({ type: 'FETCH_RESOURCE' });
-      try {
-        const response = await readResource({
-          body: {
-            session_id: sessionId,
-            uri: resourceUri,
-            extension_name: extensionName,
-          },
-        });
 
-        if (response.data) {
-          const content = response.data;
-          const rawMeta = content._meta as
-            | {
-                ui?: {
-                  csp?: McpUiResourceCsp;
-                  permissions?: McpUiResourcePermissions;
-                  prefersBorder?: boolean;
-                };
-              }
-            | undefined;
+      for (let attempt = 0; attempt <= MAX_RETRIES; attempt++) {
+        if (cancelled) return;
 
-          dispatch({
-            type: 'RESOURCE_LOADED',
-            html: content.text ?? cachedHtml ?? null,
-            meta: {
+        try {
+          const response = await readResource({
+            body: {
+              session_id: sessionId,
+              uri: resourceUri,
+              extension_name: extensionName,
+            },
+          });
+
+          if (cancelled) return;
+
+          if (response.data) {
+            const content = response.data;
+            const rawMeta = content._meta as
+              | {
+                  ui?: {
+                    csp?: McpUiResourceCsp;
+                    permissions?: McpUiResourcePermissions;
+                    prefersBorder?: boolean;
+                  };
+                }
+              | undefined;
+
+            const resolvedHtml = content.text ?? cachedHtml ?? null;
+            const resolvedMeta = {
               csp: rawMeta?.ui?.csp || null,
               // todo: pass permissions to SDK once it supports sendSandboxResourceReady
               // https://github.com/MCP-UI-Org/mcp-ui/issues/180
               permissions: null,
               prefersBorder: rawMeta?.ui?.prefersBorder ?? true,
-            },
+            };
+
+            if (resolvedHtml) {
+              fetchedDataRef.current = { html: resolvedHtml, meta: resolvedMeta };
+            }
+            dispatch({ type: 'RESOURCE_LOADED', html: resolvedHtml, meta: resolvedMeta });
+            return;
+          }
+        } catch (err) {
+          if (cancelled) return;
+
+          const isLastAttempt = attempt === MAX_RETRIES;
+
+          if (!isLastAttempt && !cachedHtml) {
+            const delay = BASE_DELAY_MS * Math.pow(2, attempt);
+            console.warn(
+              `[McpAppRenderer] Resource fetch attempt ${attempt + 1}/${MAX_RETRIES + 1} failed, retrying in ${delay}ms:`,
+              err
+            );
+            await new Promise((resolve) => setTimeout(resolve, delay));
+            continue;
+          }
+
+          console.error('[McpAppRenderer] Error fetching resource:', err);
+          if (cachedHtml) {
+            console.warn('Failed to fetch fresh resource, using cached version:', err);
+          }
+          dispatch({
+            type: 'RESOURCE_FAILED',
+            message: errorMessage(err, 'Failed to load resource'),
           });
+          return;
         }
-      } catch (err) {
-        console.error('[McpAppRenderer] Error fetching resource:', err);
-        if (cachedHtml) {
-          console.warn('Failed to fetch fresh resource, using cached version:', err);
-        }
-        dispatch({
-          type: 'RESOURCE_FAILED',
-          message: errorMessage(err, 'Failed to load resource'),
-        });
       }
     };
 
     fetchResourceData();
+
+    return () => {
+      cancelled = true;
+    };
   }, [resourceUri, extensionName, sessionId, cachedHtml]);
 
   // Create the sandbox proxy URL once we have HTML and metadata.
-  // Fetched only once — recreating the proxy would destroy iframe state.
+  // On StrictMode remount, reuse the cached URL to avoid recreating the proxy
+  // (which would destroy iframe state and cause a visible flicker).
   const pendingCsp = state.status === 'loading_sandbox' ? state.meta.csp : null;
   useEffect(() => {
     if (state.status !== 'loading_sandbox') return;
 
+    if (sandboxUrlRef.current) {
+      const { url, csp } = sandboxUrlRef.current;
+      dispatch({ type: 'SANDBOX_READY', sandboxUrl: url, sandboxCsp: csp });
+      return;
+    }
+
     fetchMcpAppProxyUrl(pendingCsp).then((url) => {
       if (url) {
+        sandboxUrlRef.current = { url, csp: pendingCsp };
         dispatch({ type: 'SANDBOX_READY', sandboxUrl: url, sandboxCsp: pendingCsp });
       } else {
         dispatch({ type: 'SANDBOX_FAILED', message: 'Failed to initialize sandbox proxy' });
@@ -458,16 +537,38 @@ export default function McpAppRenderer({
 
   const handleFallbackRequest = useCallback(
     async (request: JSONRPCRequest, _extra: RequestHandlerExtra) => {
-      // todo: handle `sampling/createMessage` per https://github.com/block/goose/pull/7039
       if (request.method === 'sampling/createMessage') {
-        return { status: 'success' as const };
+        if (!sessionId || !apiHost || !secretKey) {
+          throw new Error('Session not initialized for sampling request');
+        }
+        const { messages, systemPrompt, maxTokens } =
+          request.params as unknown as SamplingCreateMessageParams;
+        const response = await fetch(`${apiHost}/sessions/${sessionId}/sampling/message`, {
+          method: 'POST',
+          headers: {
+            'Content-Type': 'application/json',
+            'X-Secret-Key': secretKey,
+          },
+          body: JSON.stringify({
+            messages: messages.map((m) => ({
+              role: m.role,
+              content: m.content,
+            })),
+            systemPrompt,
+            maxTokens,
+          }),
+        });
+        if (!response.ok) {
+          throw new Error(`Sampling request failed: ${response.statusText}`);
+        }
+        return (await response.json()) as SamplingCreateMessageResponse;
       }
       return {
         status: 'error' as const,
         message: `Unhandled JSON-RPC method: ${request.method ?? '<unknown>'}`,
       };
     },
-    []
+    [sessionId, apiHost, secretKey]
   );
 
   const handleError = useCallback((err: Error) => {

ui/desktop/src/components/McpApps/types.ts

@@ -41,3 +41,21 @@ export type McpAppToolResult = {
   content: Content[];
   structuredContent?: unknown;
 };
+
+export type SamplingMessage = {
+  role: 'user' | 'assistant';
+  content: { type: 'text'; text: string } | { type: 'image'; data: string; mimeType: string };
+};
+
+export type SamplingCreateMessageParams = {
+  messages: SamplingMessage[];
+  systemPrompt?: string;
+  maxTokens?: number;
+};
+
+export type SamplingCreateMessageResponse = {
+  model: string;
+  stopReason: string;
+  role: 'assistant';
+  content: { type: 'text'; text: string };
+};