tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 5 Packages Projects Releases Wiki Activity

[goose2] MCP Apps: hydrate and replay app payloads in Goose2 (#8632)

Browse Source 
Signed-off-by: Andrew Harvard <aharvard@squareup.com>
This commit is contained in:
Andrew Harvard
2026-04-28 11:03:40 -04:00
committed by GitHub
parent e95811966a
commit b59ef3a46f
24 changed files with 1287 additions and 225 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/agents/extension_manager.rs
+254 -23
View File
@@ -45,8 +45,8 @@ use crate::oauth::oauth_flow;
use crate::prompt_template;
use crate::subprocess::configure_subprocess;
use rmcp::model::{
CallToolRequestParams, Content, ErrorCode, ErrorData, GetPromptResult, Prompt, Resource,
ResourceContents, ServerInfo, Tool,
CallToolRequestParams, CallToolResult, Content, ErrorCode, ErrorData, GetPromptResult, Meta,
Prompt, Resource, ResourceContents, ServerInfo, Tool,
};
use rmcp::transport::auth::AuthClient;
use schemars::_private::NoSerialize;
@@ -121,6 +121,22 @@ pub struct ExtensionManagerCapabilities {
pub host_info: Option<GooseMcpHostInfo>,
}
#[derive(Debug, Clone, serde::Serialize)]
#[serde(rename_all = "camelCase")]
pub struct GooseMcpAppToolAttachment {
pub tool_name: String,
pub extension_name: String,
pub resource_uri: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub tool_meta: Option<Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub resource_result: Option<Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub read_error: Option<String>,
}
pub(crate) const TRUSTED_TOOL_UPDATE_META_KEY: &str = "__goose_tool_update_meta";
/// Manages goose extensions / MCP clients and their interactions
pub struct ExtensionManager {
extensions: Mutex<HashMap<String, Extension>>,
@@ -214,6 +230,68 @@ pub fn get_tool_owner(tool: &Tool) -> Option<String> {
.map(|s| s.to_string())
}
fn get_tool_meta_value(tool: &Tool) -> Option<Value> {
tool.meta.as_ref().map(|meta| Value::Object(meta.0.clone()))
}
fn get_tool_resource_uri(tool: &Tool) -> Option<String> {
tool.meta
.as_ref()
.and_then(|meta| meta.0.get("ui"))
.and_then(Value::as_object)
.and_then(|ui| ui.get("resourceUri"))
.and_then(Value::as_str)
.map(ToString::to_string)
}
fn remove_untrusted_mcp_app_meta(result: &mut CallToolResult) {
let Some(meta) = result.meta.as_mut() else {
return;
};
meta.0.remove(TRUSTED_TOOL_UPDATE_META_KEY);
let remove_goose = meta
.0
.get_mut("goose")
.and_then(Value::as_object_mut)
.map(|goose_meta| {
goose_meta.remove("mcpApp");
goose_meta.is_empty()
})
.unwrap_or(false);
if remove_goose {
meta.0.remove("goose");
}
if meta.0.is_empty() {
result.meta = None;
}
}
fn insert_trusted_tool_update_meta(
result: &mut CallToolResult,
attachment: &GooseMcpAppToolAttachment,
) {
let Ok(attachment_value) = serde_json::to_value(attachment) else {
return;
};
let mut meta_map = result
.meta
.as_ref()
.map(|meta| meta.0.clone())
.unwrap_or_default();
let mut trusted_meta = serde_json::Map::new();
trusted_meta.insert("mcpApp".to_string(), attachment_value);
meta_map.insert(
TRUSTED_TOOL_UPDATE_META_KEY.to_string(),
Value::Object(trusted_meta),
);
result.meta = Some(Meta(meta_map));
}
fn is_unprefixed_extension(config: &ExtensionConfig) -> bool {
match config {
ExtensionConfig::Platform { name, .. } | ExtensionConfig::Builtin { name, .. } => {
@@ -241,9 +319,12 @@ pub fn is_hidden_extension(name: &str) -> bool {
/// Result of resolving a tool call to its owning extension
struct ResolvedTool {
tool_name: String,
extension_name: String,
actual_tool_name: String,
client: McpClientBox,
tool_meta: Option<Value>,
resource_uri: Option<String>,
}
async fn child_process_client(
@@ -1063,6 +1144,48 @@ impl ExtensionManager {
Ok(tools)
}
fn host_supports_mcp_apps(&self) -> bool {
if let Some(host_info) = &self.capabilities.host_info {
if host_info.explicit_extensions {
return host_info.mcpui_enabled();
}
}
self.capabilities.mcpui
}
async fn hydrate_mcp_app_attachment(
client: &McpClientBox,
session_id: &str,
resolved_tool: &ResolvedTool,
cancellation_token: CancellationToken,
) -> Option<GooseMcpAppToolAttachment> {
let resource_uri = resolved_tool.resource_uri.clone()?;
let mut attachment = GooseMcpAppToolAttachment {
tool_name: resolved_tool.tool_name.clone(),
extension_name: resolved_tool.extension_name.clone(),
resource_uri: resource_uri.clone(),
tool_meta: resolved_tool.tool_meta.clone(),
resource_result: None,
read_error: None,
};
match client
.read_resource(session_id, &resource_uri, cancellation_token)
.await
{
Ok(resource_result) => {
attachment.resource_result = serde_json::to_value(&resource_result).ok();
}
Err(error) => {
attachment.read_error = Some(error.to_string());
}
}
Some(attachment)
}
async fn invalidate_tools_cache_and_bump_version(&self) {
self.tools_cache_version.fetch_add(1, Ordering::SeqCst);
*self.tools_cache.lock().await = None;
@@ -1432,17 +1555,6 @@ impl ExtensionManager {
session_id: &str,
tool_name: &str,
) -> Result<ResolvedTool, ErrorData> {
if let Some((prefix, actual)) = tool_name.split_once("__") {
let owner = name_to_key(prefix);
if let Some(client) = self.get_server_client(&owner).await {
return Ok(ResolvedTool {
extension_name: owner,
actual_tool_name: actual.to_string(),
client,
});
}
}
let tools = self.get_all_tools_cached(session_id).await.map_err(|e| {
ErrorData::new(
ErrorCode::INTERNAL_ERROR,
@@ -1452,13 +1564,19 @@ impl ExtensionManager {
})?;
if let Some(tool) = tools.iter().find(|t| *t.name == *tool_name) {
let owner = get_tool_owner(tool).ok_or_else(|| {
ErrorData::new(
ErrorCode::RESOURCE_NOT_FOUND,
format!("Tool '{}' has no owner", tool_name),
None,
)
})?;
let owner = get_tool_owner(tool)
.or_else(|| {
tool_name
.split_once("__")
.map(|(prefix, _)| name_to_key(prefix))
})
.ok_or_else(|| {
ErrorData::new(
ErrorCode::RESOURCE_NOT_FOUND,
format!("Tool '{}' has no owner", tool_name),
None,
)
})?;
let actual_tool_name = tool_name
.strip_prefix(&format!("{owner}__"))
@@ -1474,12 +1592,29 @@ impl ExtensionManager {
})?;
return Ok(ResolvedTool {
tool_name: tool.name.to_string(),
extension_name: owner,
actual_tool_name,
client,
tool_meta: get_tool_meta_value(tool),
resource_uri: get_tool_resource_uri(tool),
});
}
if let Some((prefix, actual)) = tool_name.split_once("__") {
let owner = name_to_key(prefix);
if let Some(client) = self.get_server_client(&owner).await {
return Ok(ResolvedTool {
tool_name: tool_name.to_string(),
extension_name: owner,
actual_tool_name: actual.to_string(),
client,
tool_meta: None,
resource_uri: None,
});
}
}
Err(ErrorData::new(
ErrorCode::RESOURCE_NOT_FOUND,
format!("Tool '{}' not found", tool_name),
@@ -1515,8 +1650,13 @@ impl ExtensionManager {
let arguments = tool_call.arguments.clone();
let client = resolved.client.clone();
let hydration_client = client.clone();
let notifications_receiver = client.subscribe().await;
let actual_tool_name = resolved.actual_tool_name;
let actual_tool_name = resolved.actual_tool_name.clone();
let resolved_tool = resolved;
let should_hydrate_mcp_app = self.host_supports_mcp_apps();
let read_cancellation_token = cancellation_token.clone();
let session_id = ctx.session_id.clone();
let owned_ctx = ToolCallContext::new(
ctx.session_id.clone(),
ctx.working_dir.clone(),
@@ -1530,7 +1670,7 @@ impl ExtensionManager {
owned_ctx.session_id,
owned_ctx.working_dir,
);
client
let mut result = client
.call_tool(&owned_ctx, &actual_tool_name, arguments, cancellation_token)
.await
.map_err(|e| match e {
@@ -1538,7 +1678,24 @@ impl ExtensionManager {
_ => {
ErrorData::new(ErrorCode::INTERNAL_ERROR, e.to_string(), e.maybe_to_value())
}
})
})?;
remove_untrusted_mcp_app_meta(&mut result);
if should_hydrate_mcp_app && result.is_error != Some(true) {
if let Some(attachment) = Self::hydrate_mcp_app_attachment(
&hydration_client,
&session_id,
&resolved_tool,
read_cancellation_token,
)
.await
{
insert_trusted_tool_update_meta(&mut result, &attachment);
}
}
Ok(result)
};
Ok(ToolCallResult {
@@ -2318,6 +2475,80 @@ mod tests {
assert!(!tool_names.iter().any(|n| n.starts_with("ext_b__")));
}
#[test]
fn test_remove_untrusted_mcp_app_meta_strips_spoofed_payload() {
let mut result = CallToolResult::success(vec![]);
result.meta = Some(Meta(
serde_json::from_value(serde_json::json!({
"goose": {
"mcpApp": {
"resourceUri": "ui://spoofed/app",
},
"other": true,
},
TRUSTED_TOOL_UPDATE_META_KEY: {
"mcpApp": {
"resourceUri": "ui://spoofed/internal",
},
},
}))
.unwrap(),
));
remove_untrusted_mcp_app_meta(&mut result);
let meta = result.meta.expect("expected remaining meta");
assert_eq!(meta.0.get(TRUSTED_TOOL_UPDATE_META_KEY), None);
assert_eq!(
meta.0.get("goose"),
Some(&serde_json::json!({ "other": true }))
);
}
#[test]
fn test_insert_trusted_tool_update_meta_stores_backend_payload() {
let mut result = CallToolResult::success(vec![]);
let attachment = GooseMcpAppToolAttachment {
tool_name: "weather__render".to_string(),
extension_name: "weather".to_string(),
resource_uri: "ui://weather/app".to_string(),
tool_meta: None,
resource_result: Some(serde_json::json!({
"contents": [
{
"uri": "ui://weather/app",
"mimeType": "text/html;profile=mcp-app",
"text": "<div>Hello</div>",
},
],
})),
read_error: None,
};
insert_trusted_tool_update_meta(&mut result, &attachment);
let meta = result.meta.expect("expected trusted meta");
assert_eq!(
meta.0.get(TRUSTED_TOOL_UPDATE_META_KEY),
Some(&serde_json::json!({
"mcpApp": {
"toolName": "weather__render",
"extensionName": "weather",
"resourceUri": "ui://weather/app",
"resourceResult": {
"contents": [
{
"uri": "ui://weather/app",
"mimeType": "text/html;profile=mcp-app",
"text": "<div>Hello</div>",
},
],
},
},
})),
);
}
#[tokio::test]
async fn test_add_extension_noop_on_identical_config() {
// When add_extension is called with a config that is byte-for-byte identical to