Skip to content

Commit bce545c

Browse files
committed
hydrate and replay mcp app payloads in goose2
Signed-off-by: Andrew Harvard <aharvard@squareup.com>
1 parent f690c32 commit bce545c

17 files changed

Lines changed: 786 additions & 104 deletions

File tree

crates/goose-acp/src/server.rs

Lines changed: 17 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1141,12 +1141,11 @@ impl GooseAcpAgent {
11411141
}
11421142
}
11431143

1144+
let update = ToolCallUpdate::new(ToolCallId::new(tool_response.id.clone()), fields)
1145+
.meta(extract_tool_call_update_meta(tool_response));
11441146
cx.send_notification(SessionNotification::new(
11451147
session_id.clone(),
1146-
SessionUpdate::ToolCallUpdate(ToolCallUpdate::new(
1147-
ToolCallId::new(tool_response.id.clone()),
1148-
fields,
1149-
)),
1148+
SessionUpdate::ToolCallUpdate(update),
11501149
))?;
11511150

11521151
Ok(())
@@ -1238,6 +1237,16 @@ fn outcome_to_confirmation(outcome: &RequestPermissionOutcome) -> PermissionConf
12381237
}
12391238
}
12401239

1240+
fn extract_tool_call_update_meta(
1241+
tool_response: &goose::conversation::message::ToolResponse,
1242+
) -> Option<Meta> {
1243+
let tool_result = tool_response.tool_result.as_ref().ok()?;
1244+
let goose_meta = tool_result.meta.as_ref()?.0.get("goose")?.clone();
1245+
let mut meta_map = serde_json::Map::new();
1246+
meta_map.insert("goose".to_string(), goose_meta);
1247+
Some(meta_map)
1248+
}
1249+
12411250
fn build_tool_call_content(tool_result: &ToolResult<CallToolResult>) -> Vec<ToolCallContent> {
12421251
match tool_result {
12431252
Ok(result) => result
@@ -1654,12 +1663,12 @@ impl GooseAcpAgent {
16541663
}
16551664
}
16561665

1666+
let update =
1667+
ToolCallUpdate::new(ToolCallId::new(tool_response.id.clone()), fields)
1668+
.meta(extract_tool_call_update_meta(tool_response));
16571669
cx.send_notification(SessionNotification::new(
16581670
args.session_id.clone(),
1659-
SessionUpdate::ToolCallUpdate(ToolCallUpdate::new(
1660-
ToolCallId::new(tool_response.id.clone()),
1661-
fields,
1662-
)),
1671+
SessionUpdate::ToolCallUpdate(update),
16631672
))?;
16641673
}
16651674
MessageContent::Thinking(thinking) => {

crates/goose/src/agents/extension_manager.rs

Lines changed: 155 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -45,8 +45,8 @@ use crate::oauth::oauth_flow;
4545
use crate::prompt_template;
4646
use crate::subprocess::configure_subprocess;
4747
use rmcp::model::{
48-
CallToolRequestParams, Content, ErrorCode, ErrorData, GetPromptResult, Prompt, Resource,
49-
ResourceContents, ServerInfo, Tool,
48+
CallToolRequestParams, CallToolResult, Content, ErrorCode, ErrorData, GetPromptResult, Meta,
49+
Prompt, Resource, ResourceContents, ServerInfo, Tool,
5050
};
5151
use rmcp::transport::auth::AuthClient;
5252
use schemars::_private::NoSerialize;
@@ -121,6 +121,20 @@ pub struct ExtensionManagerCapabilities {
121121
pub host_info: Option<GooseMcpHostInfo>,
122122
}
123123

124+
#[derive(Debug, Clone, serde::Serialize)]
125+
#[serde(rename_all = "camelCase")]
126+
pub struct GooseMcpAppToolAttachment {
127+
pub tool_name: String,
128+
pub extension_name: String,
129+
pub resource_uri: String,
130+
#[serde(skip_serializing_if = "Option::is_none")]
131+
pub tool_meta: Option<Value>,
132+
#[serde(skip_serializing_if = "Option::is_none")]
133+
pub resource_result: Option<Value>,
134+
#[serde(skip_serializing_if = "Option::is_none")]
135+
pub read_error: Option<String>,
136+
}
137+
124138
/// Manages goose extensions / MCP clients and their interactions
125139
pub struct ExtensionManager {
126140
extensions: Mutex<HashMap<String, Extension>>,
@@ -214,6 +228,42 @@ pub fn get_tool_owner(tool: &Tool) -> Option<String> {
214228
.map(|s| s.to_string())
215229
}
216230

231+
fn get_tool_meta_value(tool: &Tool) -> Option<Value> {
232+
tool.meta.as_ref().map(|meta| Value::Object(meta.0.clone()))
233+
}
234+
235+
fn get_tool_resource_uri(tool: &Tool) -> Option<String> {
236+
tool.meta
237+
.as_ref()
238+
.and_then(|meta| meta.0.get("ui"))
239+
.and_then(Value::as_object)
240+
.and_then(|ui| ui.get("resourceUri"))
241+
.and_then(Value::as_str)
242+
.map(ToString::to_string)
243+
}
244+
245+
fn merge_mcp_app_attachment_meta(
246+
result: &mut CallToolResult,
247+
attachment: &GooseMcpAppToolAttachment,
248+
) {
249+
let mut meta_map = result
250+
.meta
251+
.as_ref()
252+
.map(|meta| meta.0.clone())
253+
.unwrap_or_default();
254+
let mut goose_value = meta_map
255+
.remove("goose")
256+
.and_then(|value| value.as_object().cloned())
257+
.unwrap_or_default();
258+
259+
goose_value.insert(
260+
"mcpApp".to_string(),
261+
serde_json::to_value(attachment).unwrap_or(Value::Null),
262+
);
263+
meta_map.insert("goose".to_string(), Value::Object(goose_value));
264+
result.meta = Some(Meta(meta_map));
265+
}
266+
217267
fn is_unprefixed_extension(config: &ExtensionConfig) -> bool {
218268
match config {
219269
ExtensionConfig::Platform { name, .. } | ExtensionConfig::Builtin { name, .. } => {
@@ -241,9 +291,12 @@ pub fn is_hidden_extension(name: &str) -> bool {
241291

242292
/// Result of resolving a tool call to its owning extension
243293
struct ResolvedTool {
294+
tool_name: String,
244295
extension_name: String,
245296
actual_tool_name: String,
246297
client: McpClientBox,
298+
tool_meta: Option<Value>,
299+
resource_uri: Option<String>,
247300
}
248301

249302
async fn child_process_client(
@@ -969,6 +1022,55 @@ impl ExtensionManager {
9691022
Ok(tools)
9701023
}
9711024

1025+
fn host_supports_mcp_apps(&self) -> bool {
1026+
if let Some(host_info) = &self.capabilities.host_info {
1027+
if host_info.explicit_extensions {
1028+
return host_info.mcpui_enabled();
1029+
}
1030+
}
1031+
1032+
self.capabilities.mcpui
1033+
}
1034+
1035+
async fn hydrate_mcp_app_attachment(
1036+
client: &McpClientBox,
1037+
session_id: &str,
1038+
resolved_tool: &ResolvedTool,
1039+
cancellation_token: CancellationToken,
1040+
result: &mut CallToolResult,
1041+
) {
1042+
if result.is_error == Some(true) {
1043+
return;
1044+
}
1045+
1046+
let Some(resource_uri) = resolved_tool.resource_uri.clone() else {
1047+
return;
1048+
};
1049+
1050+
let mut attachment = GooseMcpAppToolAttachment {
1051+
tool_name: resolved_tool.tool_name.clone(),
1052+
extension_name: resolved_tool.extension_name.clone(),
1053+
resource_uri: resource_uri.clone(),
1054+
tool_meta: resolved_tool.tool_meta.clone(),
1055+
resource_result: None,
1056+
read_error: None,
1057+
};
1058+
1059+
match client
1060+
.read_resource(session_id, &resource_uri, cancellation_token)
1061+
.await
1062+
{
1063+
Ok(resource_result) => {
1064+
attachment.resource_result = serde_json::to_value(&resource_result).ok();
1065+
}
1066+
Err(error) => {
1067+
attachment.read_error = Some(error.to_string());
1068+
}
1069+
}
1070+
1071+
merge_mcp_app_attachment_meta(result, &attachment);
1072+
}
1073+
9721074
async fn invalidate_tools_cache_and_bump_version(&self) {
9731075
self.tools_cache_version.fetch_add(1, Ordering::SeqCst);
9741076
*self.tools_cache.lock().await = None;
@@ -1338,17 +1440,6 @@ impl ExtensionManager {
13381440
session_id: &str,
13391441
tool_name: &str,
13401442
) -> Result<ResolvedTool, ErrorData> {
1341-
if let Some((prefix, actual)) = tool_name.split_once("__") {
1342-
let owner = name_to_key(prefix);
1343-
if let Some(client) = self.get_server_client(&owner).await {
1344-
return Ok(ResolvedTool {
1345-
extension_name: owner,
1346-
actual_tool_name: actual.to_string(),
1347-
client,
1348-
});
1349-
}
1350-
}
1351-
13521443
let tools = self.get_all_tools_cached(session_id).await.map_err(|e| {
13531444
ErrorData::new(
13541445
ErrorCode::INTERNAL_ERROR,
@@ -1358,13 +1449,19 @@ impl ExtensionManager {
13581449
})?;
13591450

13601451
if let Some(tool) = tools.iter().find(|t| *t.name == *tool_name) {
1361-
let owner = get_tool_owner(tool).ok_or_else(|| {
1362-
ErrorData::new(
1363-
ErrorCode::RESOURCE_NOT_FOUND,
1364-
format!("Tool '{}' has no owner", tool_name),
1365-
None,
1366-
)
1367-
})?;
1452+
let owner = get_tool_owner(tool)
1453+
.or_else(|| {
1454+
tool_name
1455+
.split_once("__")
1456+
.map(|(prefix, _)| name_to_key(prefix))
1457+
})
1458+
.ok_or_else(|| {
1459+
ErrorData::new(
1460+
ErrorCode::RESOURCE_NOT_FOUND,
1461+
format!("Tool '{}' has no owner", tool_name),
1462+
None,
1463+
)
1464+
})?;
13681465

13691466
let actual_tool_name = tool_name
13701467
.strip_prefix(&format!("{owner}__"))
@@ -1380,12 +1477,29 @@ impl ExtensionManager {
13801477
})?;
13811478

13821479
return Ok(ResolvedTool {
1480+
tool_name: tool.name.to_string(),
13831481
extension_name: owner,
13841482
actual_tool_name,
13851483
client,
1484+
tool_meta: get_tool_meta_value(tool),
1485+
resource_uri: get_tool_resource_uri(tool),
13861486
});
13871487
}
13881488

1489+
if let Some((prefix, actual)) = tool_name.split_once("__") {
1490+
let owner = name_to_key(prefix);
1491+
if let Some(client) = self.get_server_client(&owner).await {
1492+
return Ok(ResolvedTool {
1493+
tool_name: tool_name.to_string(),
1494+
extension_name: owner,
1495+
actual_tool_name: actual.to_string(),
1496+
client,
1497+
tool_meta: None,
1498+
resource_uri: None,
1499+
});
1500+
}
1501+
}
1502+
13891503
Err(ErrorData::new(
13901504
ErrorCode::RESOURCE_NOT_FOUND,
13911505
format!("Tool '{}' not found", tool_name),
@@ -1421,8 +1535,13 @@ impl ExtensionManager {
14211535

14221536
let arguments = tool_call.arguments.clone();
14231537
let client = resolved.client.clone();
1538+
let hydration_client = client.clone();
14241539
let notifications_receiver = client.subscribe().await;
1425-
let actual_tool_name = resolved.actual_tool_name;
1540+
let actual_tool_name = resolved.actual_tool_name.clone();
1541+
let resolved_tool = resolved;
1542+
let should_hydrate_mcp_app = self.host_supports_mcp_apps();
1543+
let read_cancellation_token = cancellation_token.clone();
1544+
let session_id = ctx.session_id.clone();
14261545
let owned_ctx = ToolCallContext::new(
14271546
ctx.session_id.clone(),
14281547
ctx.working_dir.clone(),
@@ -1436,15 +1555,28 @@ impl ExtensionManager {
14361555
owned_ctx.session_id,
14371556
owned_ctx.working_dir,
14381557
);
1439-
client
1558+
let mut result = client
14401559
.call_tool(&owned_ctx, &actual_tool_name, arguments, cancellation_token)
14411560
.await
14421561
.map_err(|e| match e {
14431562
ServiceError::McpError(error_data) => error_data,
14441563
_ => {
14451564
ErrorData::new(ErrorCode::INTERNAL_ERROR, e.to_string(), e.maybe_to_value())
14461565
}
1447-
})
1566+
})?;
1567+
1568+
if should_hydrate_mcp_app {
1569+
Self::hydrate_mcp_app_attachment(
1570+
&hydration_client,
1571+
&session_id,
1572+
&resolved_tool,
1573+
read_cancellation_token,
1574+
&mut result,
1575+
)
1576+
.await;
1577+
}
1578+
1579+
Ok(result)
14481580
};
14491581

14501582
Ok(ToolCallResult {

ui/goose2/src/features/chat/hooks/replayBuffer.ts

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,10 @@ export function getBufferedMessage(
3535
return replayBuffers.get(sessionId)?.find((m) => m.id === messageId);
3636
}
3737

38+
export function getReplayBuffer(sessionId: string): Message[] | undefined {
39+
return replayBuffers.get(sessionId);
40+
}
41+
3842
export function getAndDeleteReplayBuffer(
3943
sessionId: string,
4044
): Message[] | undefined {
Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
import { CodeBlock } from "@/shared/ui/ai-elements/code-block";
2+
import { useTranslation } from "react-i18next";
3+
import type { McpAppPayload } from "@/shared/types/messages";
4+
5+
interface McpAppViewProps {
6+
payload: McpAppPayload;
7+
}
8+
9+
export function McpAppView({ payload }: McpAppViewProps) {
10+
const { t } = useTranslation("chat");
11+
12+
// Currently we just render the MCP App payload as JSON.
13+
// Up next, we'll replace this with actual HTML rendering and host bridging.
14+
return (
15+
<div className="my-3" data-testid="mcp-app-view">
16+
<div className="mb-2 text-muted-foreground text-xs uppercase tracking-wide">
17+
{t("message.mcpAppUnderConstruction")}
18+
</div>
19+
<CodeBlock code={JSON.stringify(payload, null, 2)} language="json" />
20+
</div>
21+
);
22+
}

ui/goose2/src/features/chat/ui/MessageBubble.tsx

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,13 +32,15 @@ import {
3232
} from "@/shared/ui/ai-elements/reasoning";
3333
import { ToolChainCards, type ToolChainItem } from "./ToolChainCards";
3434
import { ClickableImage } from "./ClickableImage";
35+
import { McpAppView } from "./McpAppView";
3536
import { useArtifactLinkHandler } from "@/features/chat/hooks/useArtifactLinkHandler";
3637
import type {
3738
Message,
3839
MessageAttachment,
3940
MessageContent,
4041
TextContent,
4142
ImageContent,
43+
McpAppContent,
4244
ToolResponseContent,
4345
ThinkingContent,
4446
ReasoningContent as ReasoningContentType,
@@ -224,6 +226,10 @@ function renderContentBlock(
224226
case "toolResponse":
225227
// Handled by groupContentSections toolChain rendering
226228
return null;
229+
case "mcpApp": {
230+
const mcpApp = content as McpAppContent;
231+
return <McpAppView key={`mcp-app-${index}`} payload={mcpApp.payload} />;
232+
}
227233
case "thinking": {
228234
const th = content as ThinkingContent;
229235
return (

0 commit comments

Comments
 (0)