From dc56998c0f8334ab94dc39c10cfab8f0348d55fc Mon Sep 17 00:00:00 2001 From: Oliver Azevedo Barnes Date: Thu, 12 Feb 2026 13:51:04 +0000 Subject: [PATCH] agent_ui: Fix MCP tool results not displaying after app restart (#47654) Closes #47404 Release Notes: - Fixed MCP tool results not displaying after restarting Zed --- crates/agent/src/tests/mod.rs | 182 ++++++++++++++++++++++++++++++++++ crates/agent/src/thread.rs | 42 +++++--- 2 files changed, 209 insertions(+), 15 deletions(-) diff --git a/crates/agent/src/tests/mod.rs b/crates/agent/src/tests/mod.rs index 5935824b18d0095448a902c763feed3448f9fb81..499d9bdd30d50b236023e872805acaf6680f75ee 100644 --- a/crates/agent/src/tests/mod.rs +++ b/crates/agent/src/tests/mod.rs @@ -1446,6 +1446,188 @@ async fn test_mcp_tools(cx: &mut TestAppContext) { events.collect::>().await; } +#[gpui::test] +async fn test_mcp_tool_result_displayed_when_server_disconnected(cx: &mut TestAppContext) { + let ThreadTest { + model, + thread, + context_server_store, + fs, + .. + } = setup(cx, TestModel::Fake).await; + let fake_model = model.as_fake(); + + // Setup settings to allow MCP tools + fs.insert_file( + paths::settings_file(), + json!({ + "agent": { + "always_allow_tool_actions": true, + "profiles": { + "test": { + "name": "Test Profile", + "enable_all_context_servers": true, + "tools": {} + }, + } + } + }) + .to_string() + .into_bytes(), + ) + .await; + cx.run_until_parked(); + thread.update(cx, |thread, cx| { + thread.set_profile(AgentProfileId("test".into()), cx) + }); + + // Setup a context server with a tool + let mut mcp_tool_calls = setup_context_server( + "github_server", + vec![context_server::types::Tool { + name: "issue_read".into(), + description: Some("Read a GitHub issue".into()), + input_schema: json!({ + "type": "object", + "properties": { + "issue_url": { "type": "string" } + } + }), + output_schema: None, + annotations: None, + }], + &context_server_store, + cx, + ); + + // Send a message and have the model call the MCP tool + let events = thread.update(cx, |thread, cx| { + thread + .send(UserMessageId::new(), ["Read issue #47404"], cx) + .unwrap() + }); + cx.run_until_parked(); + + // Verify the MCP tool is available to the model + let completion = fake_model.pending_completions().pop().unwrap(); + assert_eq!( + tool_names_for_completion(&completion), + vec!["issue_read"], + "MCP tool should be available" + ); + + // Simulate the model calling the MCP tool + fake_model.send_last_completion_stream_event(LanguageModelCompletionEvent::ToolUse( + LanguageModelToolUse { + id: "tool_1".into(), + name: "issue_read".into(), + raw_input: json!({"issue_url": "https://github.com/zed-industries/zed/issues/47404"}) + .to_string(), + input: json!({"issue_url": "https://github.com/zed-industries/zed/issues/47404"}), + is_input_complete: true, + thought_signature: None, + }, + )); + fake_model.end_last_completion_stream(); + cx.run_until_parked(); + + // The MCP server receives the tool call and responds with content + let expected_tool_output = "Issue #47404: Tool call results are cleared upon app close"; + let (tool_call_params, tool_call_response) = mcp_tool_calls.next().await.unwrap(); + assert_eq!(tool_call_params.name, "issue_read"); + tool_call_response + .send(context_server::types::CallToolResponse { + content: vec![context_server::types::ToolResponseContent::Text { + text: expected_tool_output.into(), + }], + is_error: None, + meta: None, + structured_content: None, + }) + .unwrap(); + cx.run_until_parked(); + + // After tool completes, the model continues with a new completion request + // that includes the tool results. We need to respond to this. + let _completion = fake_model.pending_completions().pop().unwrap(); + fake_model.send_last_completion_stream_text_chunk("I found the issue!"); + fake_model + .send_last_completion_stream_event(LanguageModelCompletionEvent::Stop(StopReason::EndTurn)); + fake_model.end_last_completion_stream(); + events.collect::>().await; + + // Verify the tool result is stored in the thread by checking the markdown output. + // The tool result is in the first assistant message (not the last one, which is + // the model's response after the tool completed). + thread.update(cx, |thread, _cx| { + let markdown = thread.to_markdown(); + assert!( + markdown.contains("**Tool Result**: issue_read"), + "Thread should contain tool result header" + ); + assert!( + markdown.contains(expected_tool_output), + "Thread should contain tool output: {}", + expected_tool_output + ); + }); + + // Simulate app restart: disconnect the MCP server. + // After restart, the MCP server won't be connected yet when the thread is replayed. + context_server_store.update(cx, |store, cx| { + let _ = store.stop_server(&ContextServerId("github_server".into()), cx); + }); + cx.run_until_parked(); + + // Replay the thread (this is what happens when loading a saved thread) + let mut replay_events = thread.update(cx, |thread, cx| thread.replay(cx)); + + let mut found_tool_call = None; + let mut found_tool_call_update_with_output = None; + + while let Some(event) = replay_events.next().await { + let event = event.unwrap(); + match &event { + ThreadEvent::ToolCall(tc) if tc.tool_call_id.to_string() == "tool_1" => { + found_tool_call = Some(tc.clone()); + } + ThreadEvent::ToolCallUpdate(acp_thread::ToolCallUpdate::UpdateFields(update)) + if update.tool_call_id.to_string() == "tool_1" => + { + if update.fields.raw_output.is_some() { + found_tool_call_update_with_output = Some(update.clone()); + } + } + _ => {} + } + } + + // The tool call should be found + assert!( + found_tool_call.is_some(), + "Tool call should be emitted during replay" + ); + + assert!( + found_tool_call_update_with_output.is_some(), + "ToolCallUpdate with raw_output should be emitted even when MCP server is disconnected." + ); + + let update = found_tool_call_update_with_output.unwrap(); + assert_eq!( + update.fields.raw_output, + Some(expected_tool_output.into()), + "raw_output should contain the saved tool result" + ); + + // Also verify the status is correct (completed, not failed) + assert_eq!( + update.fields.status, + Some(acp::ToolCallStatus::Completed), + "Tool call status should reflect the original completion status" + ); +} + #[gpui::test] async fn test_mcp_tool_truncation(cx: &mut TestAppContext) { let ThreadTest { diff --git a/crates/agent/src/thread.rs b/crates/agent/src/thread.rs index d08bc1c9186d4578e759aefe58e0fe50f7982c7f..26026175e68f3fecd20d21441bb7f1e41a438207 100644 --- a/crates/agent/src/thread.rs +++ b/crates/agent/src/thread.rs @@ -982,6 +982,20 @@ impl Thread { stream: &ThreadEventStream, cx: &mut Context, ) { + // Extract saved output and status first, so they're available even if tool is not found + let output = tool_result + .as_ref() + .and_then(|result| result.output.clone()); + let status = tool_result + .as_ref() + .map_or(acp::ToolCallStatus::Failed, |result| { + if result.is_error { + acp::ToolCallStatus::Failed + } else { + acp::ToolCallStatus::Completed + } + }); + let tool = self.tools.get(tool_use.name.as_ref()).cloned().or_else(|| { self.context_server_registry .read(cx) @@ -996,14 +1010,25 @@ impl Thread { }); let Some(tool) = tool else { + // Tool not found (e.g., MCP server not connected after restart), + // but still display the saved result if available. + // We need to send both ToolCall and ToolCallUpdate events because the UI + // only converts raw_output to displayable content in update_fields, not from_acp. stream .0 .unbounded_send(Ok(ThreadEvent::ToolCall( acp::ToolCall::new(tool_use.id.to_string(), tool_use.name.to_string()) - .status(acp::ToolCallStatus::Failed) + .status(status) .raw_input(tool_use.input.clone()), ))) .ok(); + stream.update_tool_call_fields( + &tool_use.id, + acp::ToolCallUpdateFields::new() + .status(status) + .raw_output(output), + None, + ); return; }; @@ -1017,9 +1042,6 @@ impl Thread { tool_use.input.clone(), ); - let output = tool_result - .as_ref() - .and_then(|result| result.output.clone()); if let Some(output) = output.clone() { // For replay, we use a dummy cancellation receiver since the tool already completed let (_cancellation_tx, cancellation_rx) = watch::channel(false); @@ -1036,17 +1058,7 @@ impl Thread { stream.update_tool_call_fields( &tool_use.id, acp::ToolCallUpdateFields::new() - .status( - tool_result - .as_ref() - .map_or(acp::ToolCallStatus::Failed, |result| { - if result.is_error { - acp::ToolCallStatus::Failed - } else { - acp::ToolCallStatus::Completed - } - }), - ) + .status(status) .raw_output(output), None, );