From e5b1c8a2bc79066190c0196acf44176eb7f88af0 Mon Sep 17 00:00:00 2001 From: Andrew Harvard Date: Fri, 20 Mar 2026 15:56:43 -0400 Subject: [PATCH 1/5] feat: support ui/update-model-context for MCP Apps Implement the ui/update-model-context method from the MCP Apps specification, allowing MCP App views to push context updates to the host agent without triggering immediate action. Context updates are stored in-memory (last-write-wins per app key) and injected alongside MOIM as a synthetic user message on the next agent turn, giving the model awareness of app state. - Add McpAppContext store on Agent with update/collect methods - Add /agent/update_model_context server endpoint - Integrate with MOIM injection in moim.rs - Handle ui/update-model-context in McpAppRenderer - Add 7 unit tests covering store, overwrite, collect, and injection Signed-off-by: Andrew Harvard Co-authored-by: Goose Ai-assisted: true --- crates/goose-server/src/openapi.rs | 2 + crates/goose-server/src/routes/agent.rs | 49 +++++++ crates/goose/src/agents/agent.rs | 132 ++++++++++++++++++ crates/goose/src/agents/mod.rs | 4 +- crates/goose/src/agents/moim.rs | 79 ++++++++++- ui/desktop/openapi.json | 61 ++++++++ ui/desktop/src/api/index.ts | 4 +- ui/desktop/src/api/sdk.gen.ts | 11 +- ui/desktop/src/api/types.gen.ts | 39 ++++++ .../src/components/McpApps/McpAppRenderer.tsx | 30 +++- 10 files changed, 399 insertions(+), 12 deletions(-) diff --git a/crates/goose-server/src/openapi.rs b/crates/goose-server/src/openapi.rs index eec8aeefe2ac..66e0cc95adf6 100644 --- a/crates/goose-server/src/openapi.rs +++ b/crates/goose-server/src/openapi.rs @@ -425,6 +425,7 @@ derive_utoipa!(Icon as IconSchema); super::routes::agent::get_tools, super::routes::agent::read_resource, super::routes::agent::call_tool, + super::routes::agent::update_model_context, super::routes::agent::list_apps, super::routes::agent::export_app, super::routes::agent::import_app, @@ -640,6 +641,7 @@ derive_utoipa!(Icon as IconSchema); super::routes::agent::ReadResourceResponse, super::routes::agent::CallToolRequest, super::routes::agent::CallToolResponse, + super::routes::agent::UpdateModelContextRequest, ContentBlockSchema, super::routes::agent::ListAppsRequest, super::routes::agent::ListAppsResponse, diff --git a/crates/goose-server/src/routes/agent.rs b/crates/goose-server/src/routes/agent.rs index 392a00e3db38..6229e783b85a 100644 --- a/crates/goose-server/src/routes/agent.rs +++ b/crates/goose-server/src/routes/agent.rs @@ -169,6 +169,18 @@ pub struct CallToolResponse { _meta: Option, } +#[derive(Deserialize, utoipa::ToSchema)] +#[serde(rename_all = "camelCase")] +pub struct UpdateModelContextRequest { + session_id: String, + /// Key identifying the MCP App (typically `"{extension}__{resourceUri}"`) + key: String, + #[serde(default)] + content: Option>, + #[serde(default)] + structured_content: Option, +} + #[derive(Serialize, utoipa::ToSchema)] pub struct ResumeAgentResponse { pub session: Session, @@ -1100,6 +1112,42 @@ async fn call_tool( })) } +#[utoipa::path( + post, + path = "/agent/update_model_context", + request_body = UpdateModelContextRequest, + responses( + (status = 200, description = "Model context updated"), + (status = 401, description = "Unauthorized"), + (status = 424, description = "Agent not initialized"), + (status = 500, description = "Internal server error") + ), + security( + ("api_key" = []) + ), + tag = "Agent" +)] +async fn update_model_context( + State(state): State>, + Json(payload): Json, +) -> Result<(), StatusCode> { + let agent = state + .get_agent_for_route(payload.session_id.clone()) + .await?; + + agent + .update_mcp_app_context( + payload.key, + goose::agents::McpAppContext { + content: payload.content, + structured_content: payload.structured_content, + }, + ) + .await; + + Ok(()) +} + #[derive(Deserialize, utoipa::IntoParams, utoipa::ToSchema)] pub struct ListAppsRequest { session_id: Option, @@ -1306,6 +1354,7 @@ pub fn routes(state: Arc) -> Router { .route("/agent/tools", get(get_tools)) .route("/agent/read_resource", post(read_resource)) .route("/agent/call_tool", post(call_tool)) + .route("/agent/update_model_context", post(update_model_context)) .route("/agent/list_apps", get(list_apps)) .route("/agent/export_app/{name}", get(export_app)) .route("/agent/import_app", post(import_app)) diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index ee31d401e112..87328ba2774e 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -64,6 +64,14 @@ use tracing::{debug, error, info, instrument, warn}; const DEFAULT_MAX_TURNS: u32 = 1000; const COMPACTION_THINKING_TEXT: &str = "goose is compacting the conversation..."; +/// Context provided by an MCP App via `ui/update-model-context`. +/// Each update overwrites the previous one for the same key. +#[derive(Clone, Debug)] +pub struct McpAppContext { + pub content: Option>, + pub structured_content: Option, +} + /// Context needed for the reply function pub struct ReplyContext { pub conversation: Conversation, @@ -152,6 +160,9 @@ pub struct Agent { pub(super) retry_manager: RetryManager, pub(super) tool_inspection_manager: ToolInspectionManager, container: Mutex>, + /// Context provided by MCP Apps via `ui/update-model-context`, keyed by + /// `"{extension_name}__{resource_uri}"`. Only the latest update per key is kept. + pub(super) mcp_app_contexts: Mutex>, } #[derive(Clone, Debug)] @@ -250,6 +261,7 @@ impl Agent { provider.clone(), ), container: Mutex::new(None), + mcp_app_contexts: Mutex::new(HashMap::new()), } } @@ -481,6 +493,42 @@ impl Agent { self.container.lock().await.clone() } + /// Store context from an MCP App's `ui/update-model-context` request. + pub async fn update_mcp_app_context(&self, key: String, context: McpAppContext) { + self.mcp_app_contexts.lock().await.insert(key, context); + } + + /// Collect all stored MCP App contexts for injection into the conversation. + pub async fn collect_mcp_app_contexts(&self) -> Option { + let contexts = self.mcp_app_contexts.lock().await; + if contexts.is_empty() { + return None; + } + + let mut parts = Vec::new(); + for (key, ctx) in contexts.iter() { + let mut section = format!("\n"); + if let Some(content) = &ctx.content { + for block in content { + if let Some(text) = block.get("text").and_then(|v| v.as_str()) { + section.push_str(text); + section.push('\n'); + } + } + } + if let Some(structured) = &ctx.structured_content { + if let Ok(json_str) = serde_json::to_string_pretty(structured) { + section.push_str(&json_str); + section.push('\n'); + } + } + section.push_str(""); + parts.push(section); + } + + Some(parts.join("\n")) + } + /// Check if a tool is a frontend tool pub async fn is_frontend_tool(&self, name: &str) -> bool { self.frontend_tools.lock().await.contains_key(name) @@ -1184,11 +1232,13 @@ impl Agent { break; } + let mcp_app_context = self.collect_mcp_app_contexts().await; let conversation_with_moim = super::moim::inject_moim( &session_config.id, conversation.clone(), &self.extension_manager, &working_dir, + mcp_app_context, ).await; let mut stream = Self::stream_response_from_provider( @@ -2308,4 +2358,86 @@ mod tests { Ok(()) } + + #[tokio::test] + async fn test_mcp_app_context_empty_returns_none() { + let agent = Agent::new(); + assert!(agent.collect_mcp_app_contexts().await.is_none()); + } + + #[tokio::test] + async fn test_mcp_app_context_stores_and_collects() { + let agent = Agent::new(); + + agent + .update_mcp_app_context( + "ext__ui://ext/tool".to_string(), + McpAppContext { + content: Some(vec![ + serde_json::json!({"type": "text", "text": "User selected 3 items"}), + ]), + structured_content: None, + }, + ) + .await; + + let result = agent.collect_mcp_app_contexts().await; + assert!(result.is_some()); + let text = result.unwrap(); + assert!(text.contains("User selected 3 items")); + assert!(text.contains("ext__ui://ext/tool")); + } + + #[tokio::test] + async fn test_mcp_app_context_overwrites_previous() { + let agent = Agent::new(); + let key = "ext__ui://ext/tool".to_string(); + + agent + .update_mcp_app_context( + key.clone(), + McpAppContext { + content: Some(vec![ + serde_json::json!({"type": "text", "text": "old context"}), + ]), + structured_content: None, + }, + ) + .await; + + agent + .update_mcp_app_context( + key, + McpAppContext { + content: Some(vec![ + serde_json::json!({"type": "text", "text": "new context"}), + ]), + structured_content: None, + }, + ) + .await; + + let text = agent.collect_mcp_app_contexts().await.unwrap(); + assert!(text.contains("new context")); + assert!(!text.contains("old context")); + } + + #[tokio::test] + async fn test_mcp_app_context_structured_content() { + let agent = Agent::new(); + + agent + .update_mcp_app_context( + "ext__ui://ext/tool".to_string(), + McpAppContext { + content: None, + structured_content: Some(serde_json::json!({"items": 3, "total": 150})), + }, + ) + .await; + + let text = agent.collect_mcp_app_contexts().await.unwrap(); + assert!(text.contains("\"items\": 3")); + assert!(text.contains("\"total\": 150")); + } } diff --git a/crates/goose/src/agents/mod.rs b/crates/goose/src/agents/mod.rs index 1b41a743182b..45b6c88db621 100644 --- a/crates/goose/src/agents/mod.rs +++ b/crates/goose/src/agents/mod.rs @@ -23,7 +23,9 @@ mod tool_execution; pub mod types; pub mod validate_extensions; -pub use agent::{Agent, AgentConfig, AgentEvent, ExtensionLoadResult, GoosePlatform}; +pub use agent::{ + Agent, AgentConfig, AgentEvent, ExtensionLoadResult, GoosePlatform, McpAppContext, +}; pub use container::Container; pub use execute_commands::COMPACT_TRIGGERS; pub use extension::{ExtensionConfig, ExtensionError}; diff --git a/crates/goose/src/agents/moim.rs b/crates/goose/src/agents/moim.rs index 840017fb1706..fb5e8871779c 100644 --- a/crates/goose/src/agents/moim.rs +++ b/crates/goose/src/agents/moim.rs @@ -14,21 +14,30 @@ pub async fn inject_moim( conversation: Conversation, extension_manager: &ExtensionManager, working_dir: &Path, + mcp_app_context: Option, ) -> Conversation { if SKIP.with(|f| f.get()) { return conversation; } - if let Some(moim) = extension_manager + let moim = extension_manager .collect_moim(session_id, working_dir) - .await - { + .await; + + let combined = match (moim, mcp_app_context) { + (Some(m), Some(ctx)) => Some(format!("{m}\n{ctx}")), + (Some(m), None) => Some(m), + (None, Some(ctx)) => Some(ctx), + (None, None) => None, + }; + + if let Some(injected_text) = combined { let mut messages = conversation.messages().clone(); let idx = messages .iter() .rposition(|m| m.role == Role::Assistant) .unwrap_or(0); - messages.insert(idx, Message::user().with_text(moim)); + messages.insert(idx, Message::user().with_text(injected_text)); let (fixed, issues) = fix_conversation(Conversation::new_unvalidated(messages)); @@ -65,7 +74,7 @@ mod tests { Message::assistant().with_text("Hi"), Message::user().with_text("Bye"), ]); - let result = inject_moim("test-session-id", conv, &em, &working_dir).await; + let result = inject_moim("test-session-id", conv, &em, &working_dir, None).await; let msgs = result.messages(); assert_eq!(msgs.len(), 3); @@ -90,7 +99,7 @@ mod tests { let working_dir = PathBuf::from("/test/dir"); let conv = Conversation::new_unvalidated(vec![Message::user().with_text("Hello")]); - let result = inject_moim("test-session-id", conv, &em, &working_dir).await; + let result = inject_moim("test-session-id", conv, &em, &working_dir, None).await; assert_eq!(result.messages().len(), 1); @@ -125,7 +134,7 @@ mod tests { .with_tool_response("search_2", Ok(rmcp::model::CallToolResult::success(vec![]))), ]); - let result = inject_moim("test-session-id", conv, &em, &working_dir).await; + let result = inject_moim("test-session-id", conv, &em, &working_dir, None).await; let msgs = result.messages(); assert_eq!(msgs.len(), 6); @@ -141,4 +150,60 @@ mod tests { "MOIM should be in message before latest assistant message" ); } + + #[tokio::test] + async fn test_mcp_app_context_injected_alongside_moim() { + let temp_dir = tempfile::tempdir().unwrap(); + let em = ExtensionManager::new_without_provider(temp_dir.path().to_path_buf()); + let working_dir = PathBuf::from("/test/dir"); + + let conv = Conversation::new_unvalidated(vec![ + Message::user().with_text("Hello"), + Message::assistant().with_text("Hi"), + Message::user().with_text("Bye"), + ]); + + let app_ctx = Some( + "\nUser selected 3 items\n" + .to_string(), + ); + let result = inject_moim("test-session-id", conv, &em, &working_dir, app_ctx).await; + + let all_text: String = result + .messages() + .iter() + .flat_map(|m| m.content.iter()) + .filter_map(|c| c.as_text()) + .collect::>() + .join("\n"); + + assert!(all_text.contains("\nSome context\n".to_string(), + ); + let result = inject_moim("test-session-id", conv, &em, &working_dir, app_ctx).await; + + let all_text: String = result + .messages() + .iter() + .flat_map(|m| m.content.iter()) + .filter_map(|c| c.as_text()) + .collect::>() + .join("\n"); + + assert!(all_text.contains(" = Options2 & { /** @@ -123,6 +123,15 @@ export const updateFromSession = (options: } }); +export const updateModelContext = (options: Options) => (options.client ?? client).post({ + url: '/agent/update_model_context', + ...options, + headers: { + 'Content-Type': 'application/json', + ...options.headers + } +}); + export const updateAgentProvider = (options: Options) => (options.client ?? client).post({ url: '/agent/update_provider', ...options, diff --git a/ui/desktop/src/api/types.gen.ts b/ui/desktop/src/api/types.gen.ts index 6add7592e84e..8c00b1f5fde9 100644 --- a/ui/desktop/src/api/types.gen.ts +++ b/ui/desktop/src/api/types.gen.ts @@ -1550,6 +1550,16 @@ export type UpdateFromSessionRequest = { session_id: string; }; +export type UpdateModelContextRequest = { + content?: Array | null; + /** + * Key identifying the MCP App (typically `"{extension}__{resourceUri}"`) + */ + key: string; + sessionId: string; + structuredContent?: unknown; +}; + export type UpdateProviderRequest = { context_limit?: number | null; model?: string | null; @@ -2074,6 +2084,35 @@ export type UpdateFromSessionResponses = { 200: unknown; }; +export type UpdateModelContextData = { + body: UpdateModelContextRequest; + path?: never; + query?: never; + url: '/agent/update_model_context'; +}; + +export type UpdateModelContextErrors = { + /** + * Unauthorized + */ + 401: unknown; + /** + * Agent not initialized + */ + 424: unknown; + /** + * Internal server error + */ + 500: unknown; +}; + +export type UpdateModelContextResponses = { + /** + * Model context updated + */ + 200: unknown; +}; + export type UpdateAgentProviderData = { body: UpdateProviderRequest; path?: never; diff --git a/ui/desktop/src/components/McpApps/McpAppRenderer.tsx b/ui/desktop/src/components/McpApps/McpAppRenderer.tsx index 4d442e69fb29..34d094ea51f9 100644 --- a/ui/desktop/src/components/McpApps/McpAppRenderer.tsx +++ b/ui/desktop/src/components/McpApps/McpAppRenderer.tsx @@ -599,6 +599,34 @@ export default function McpAppRenderer({ const handleFallbackRequest = useCallback( async (request: JSONRPCRequest, _extra: RequestHandlerExtra) => { + if (request.method === 'ui/update-model-context') { + if (!sessionId || !apiHost || !secretKey) { + throw new Error('Session not initialized for model context update'); + } + const params = request.params as { + content?: Array<{ type: string; text?: string }>; + structuredContent?: Record; + }; + const key = `${extensionName}__${resourceUri}`; + const response = await fetch(`${apiHost}/agent/update_model_context`, { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + 'X-Secret-Key': secretKey, + }, + body: JSON.stringify({ + sessionId, + key, + content: params?.content ?? null, + structuredContent: params?.structuredContent ?? null, + }), + }); + if (!response.ok) { + throw new Error(`Model context update failed: ${response.statusText}`); + } + return {}; + } + if (request.method === 'sampling/createMessage') { if (!sessionId || !apiHost || !secretKey) { throw new Error('Session not initialized for sampling request'); @@ -630,7 +658,7 @@ export default function McpAppRenderer({ message: `Unhandled JSON-RPC method: ${request.method ?? ''}`, }; }, - [sessionId, apiHost, secretKey] + [sessionId, apiHost, secretKey, extensionName, resourceUri] ); const handleError = useCallback((err: Error) => { From 8d191713b6212b1754e24f78fba321687ce62813 Mon Sep 17 00:00:00 2001 From: Andrew Harvard Date: Fri, 20 Mar 2026 16:13:30 -0400 Subject: [PATCH 2/5] fix: XML-escape MCP App context before injection Escape <, >, &, and " in app-provided text and keys before wrapping them in XML blocks. Prevents malformed markup from benign payloads (e.g. code snippets) and blocks malicious apps from closing the tag to spoof extra context sections. Adds a test covering special character escaping. Signed-off-by: Andrew Harvard Co-authored-by: Goose Ai-assisted: true --- crates/goose/src/agents/agent.rs | 40 ++++++++++++++++++++++++++++---- 1 file changed, 35 insertions(+), 5 deletions(-) diff --git a/crates/goose/src/agents/agent.rs b/crates/goose/src/agents/agent.rs index 87328ba2774e..0ad7bde101be 100644 --- a/crates/goose/src/agents/agent.rs +++ b/crates/goose/src/agents/agent.rs @@ -64,6 +64,13 @@ use tracing::{debug, error, info, instrument, warn}; const DEFAULT_MAX_TURNS: u32 = 1000; const COMPACTION_THINKING_TEXT: &str = "goose is compacting the conversation..."; +fn xml_escape(s: &str) -> String { + s.replace('&', "&") + .replace('<', "<") + .replace('>', ">") + .replace('"', """) +} + /// Context provided by an MCP App via `ui/update-model-context`. /// Each update overwrites the previous one for the same key. #[derive(Clone, Debug)] @@ -507,18 +514,19 @@ impl Agent { let mut parts = Vec::new(); for (key, ctx) in contexts.iter() { - let mut section = format!("\n"); + let escaped_key = xml_escape(key); + let mut section = format!("\n"); if let Some(content) = &ctx.content { for block in content { if let Some(text) = block.get("text").and_then(|v| v.as_str()) { - section.push_str(text); + section.push_str(&xml_escape(text)); section.push('\n'); } } } if let Some(structured) = &ctx.structured_content { if let Ok(json_str) = serde_json::to_string_pretty(structured) { - section.push_str(&json_str); + section.push_str(&xml_escape(&json_str)); section.push('\n'); } } @@ -2437,7 +2445,29 @@ mod tests { .await; let text = agent.collect_mcp_app_contexts().await.unwrap(); - assert!(text.contains("\"items\": 3")); - assert!(text.contains("\"total\": 150")); + assert!(text.contains(""items": 3")); + assert!(text.contains(""total": 150")); + } + + #[tokio::test] + async fn test_mcp_app_context_escapes_xml_special_chars() { + let agent = Agent::new(); + + agent + .update_mcp_app_context( + "ext__