From e2917a2b0748279c0ff3c4f5ef86c504817e344e Mon Sep 17 00:00:00 2001 From: Alex Hancock Date: Fri, 6 Feb 2026 12:34:41 +0100 Subject: [PATCH] feat(apps): add support for MCP apps to sample Co-authored-by: Andrew Harvard --- crates/goose-server/src/routes/mod.rs | 2 + crates/goose-server/src/routes/sampling.rs | 87 +++++++++ crates/goose/src/goose_apps/cache.rs | 11 +- crates/goose/src/goose_apps/chat.html | 184 ++++++++++++++++++ .../src/components/McpApps/McpAppRenderer.tsx | 183 +++++++++++++---- ui/desktop/src/components/McpApps/types.ts | 18 ++ 6 files changed, 440 insertions(+), 45 deletions(-) create mode 100644 crates/goose-server/src/routes/sampling.rs create mode 100644 crates/goose/src/goose_apps/chat.html diff --git a/crates/goose-server/src/routes/mod.rs b/crates/goose-server/src/routes/mod.rs index e0935c2476a8..9eb00449f16e 100644 --- a/crates/goose-server/src/routes/mod.rs +++ b/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, 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)) } diff --git a/crates/goose-server/src/routes/sampling.rs b/crates/goose-server/src/routes/sampling.rs new file mode 100644 index 000000000000..daecd1b11be8 --- /dev/null +++ b/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) -> Router { + Router::new() + .route( + "/sessions/{session_id}/sampling/message", + post(create_message), + ) + .with_state(state) +} + +async fn create_message( + State(state): State>, + Path(session_id): Path, + Json(request): Json, +) -> Result, 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 = 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) -> 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 +} diff --git a/crates/goose/src/goose_apps/cache.rs b/crates/goose/src/goose_apps/cache.rs index 633701ce6350..c63640304537 100644 --- a/crates/goose/src/goose_apps/cache.rs +++ b/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); + } } } } diff --git a/crates/goose/src/goose_apps/chat.html b/crates/goose/src/goose_apps/chat.html new file mode 100644 index 000000000000..906d6a8500cd --- /dev/null +++ b/crates/goose/src/goose_apps/chat.html @@ -0,0 +1,184 @@ + + + + + Chat + + + + +
+
+ + +
+ + + + diff --git a/ui/desktop/src/components/McpApps/McpAppRenderer.tsx b/ui/desktop/src/components/McpApps/McpAppRenderer.tsx index d691930c5080..cd1646fcbbd7 100644 --- a/ui/desktop/src/components/McpApps/McpAppRenderer.tsx +++ b/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(null); const [containerWidth, setContainerWidth] = useState(0); const [containerHeight, setContainerHeight] = useState(0); + const [apiHost, setApiHost] = useState(null); + const [secretKey, setSecretKey] = useState(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 ?? ''}`, }; }, - [] + [sessionId, apiHost, secretKey] ); const handleError = useCallback((err: Error) => { diff --git a/ui/desktop/src/components/McpApps/types.ts b/ui/desktop/src/components/McpApps/types.ts index b33b131c25a0..a4a86284e112 100644 --- a/ui/desktop/src/components/McpApps/types.ts +++ b/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 }; +};