feat(mcp): support sampling in a scoped way

2026-07-03 14:10:03 +02:00 · 2025-10-24 15:09:12 -04:00
parent 565dcb879b
commit fa587e7a35
13 changed files with 294 additions and 93 deletions
@@ -117,6 +117,21 @@ jobs:
          # Run the provider test script (binary already built and downloaded)
          bash scripts/test_providers.sh

+      - name: Run MCP Tests
+        env:
+          ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
+          OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
+          GOOGLE_API_KEY: ${{ secrets.GOOGLE_API_KEY }}
+          DATABRICKS_HOST: ${{ secrets.DATABRICKS_HOST }}
+          DATABRICKS_TOKEN: ${{ secrets.DATABRICKS_TOKEN }}
+          OPENROUTER_API_KEY: ${{ secrets.OPENROUTER_API_KEY }}
+          TETRATE_API_KEY: ${{ secrets.TETRATE_API_KEY }}
+          HOME: /tmp/goose-home
+          GOOSE_DISABLE_KEYRING: 1
+          SKIP_BUILD: 1
+        run: |
+          bash scripts/test_mcp.sh
+
      - name: Run Subrecipe Tests
        env:
          ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
@@ -28,7 +28,7 @@ use crate::agents::subagent_execution_tool::tasks_manager::TasksManager;
 use crate::agents::tool_route_manager::ToolRouteManager;
 use crate::agents::tool_router_index_manager::ToolRouterIndexManager;
 use crate::agents::types::SessionConfig;
-use crate::agents::types::{FrontendTool, ToolResultReceiver};
+use crate::agents::types::{FrontendTool, SharedProvider, ToolResultReceiver};
 use crate::config::{get_enabled_extensions, Config};
 use crate::context_mgmt::DEFAULT_COMPACTION_THRESHOLD;
 use crate::conversation::{debug_conversation_fix, fix_conversation, Conversation};
@@ -86,7 +86,8 @@ pub struct ToolCategorizeResult {

 /// The main goose Agent
 pub struct Agent {
-    pub(super) provider: Mutex<Option<Arc<dyn Provider>>>,
+    pub(super) provider: SharedProvider,
+
    pub extension_manager: Arc<ExtensionManager>,
    pub(super) sub_recipe_manager: Mutex<SubRecipeManager>,
    pub(super) tasks_manager: TasksManager,
@@ -159,10 +160,11 @@ impl Agent {
        // Create channels with buffer size 32 (adjust if needed)
        let (confirm_tx, confirm_rx) = mpsc::channel(32);
        let (tool_tx, tool_rx) = mpsc::channel(32);
+        let provider = Arc::new(Mutex::new(None));

        Self {
-            provider: Mutex::new(None),
-            extension_manager: Arc::new(ExtensionManager::new()),
+            provider: provider.clone(),
+            extension_manager: Arc::new(ExtensionManager::new(provider.clone())),
            sub_recipe_manager: Mutex::new(SubRecipeManager::new()),
            tasks_manager: TasksManager::new(),
            final_output_tool: Arc::new(Mutex::new(None)),
@@ -12,6 +12,7 @@ use rmcp::transport::{
    TokioChildProcess,
 };
 use std::collections::HashMap;
+use std::option::Option;
 use std::process::Stdio;
 use std::sync::Arc;
 use std::time::Duration;
@@ -29,6 +30,7 @@ use super::extension::{
    ToolInfo, PLATFORM_EXTENSIONS,
 };
 use super::tool_execution::ToolCallResult;
+use super::types::SharedProvider;
 use crate::agents::extension::{Envs, ProcessExit};
 use crate::agents::extension_malware_check;
 use crate::agents::mcp_client::{McpClient, McpClientTrait};
@@ -91,6 +93,7 @@ impl Extension {
 pub struct ExtensionManager {
    extensions: Mutex<HashMap<String, Extension>>,
    context: Mutex<PlatformExtensionContext>,
+    provider: SharedProvider,
 }

 /// A flattened representation of a resource used by the agent to prepare inference
@@ -171,13 +174,14 @@ pub fn get_parameter_names(tool: &Tool) -> Vec<String> {

 impl Default for ExtensionManager {
    fn default() -> Self {
-        Self::new()
+        Self::new(Arc::new(Mutex::new(None)))
    }
 }

 async fn child_process_client(
    mut command: Command,
    timeout: &Option<u64>,
+    provider: SharedProvider,
 ) -> ExtensionResult<McpClient> {
    #[cfg(unix)]
    command.process_group(0);
@@ -205,6 +209,7 @@ async fn child_process_client(
    let client_result = McpClient::connect(
        transport,
        Duration::from_secs(timeout.unwrap_or(crate::config::DEFAULT_EXTENSION_TIMEOUT)),
+        provider,
    )
    .await;

@@ -243,7 +248,7 @@ fn extract_auth_error(
 }

 impl ExtensionManager {
-    pub fn new() -> Self {
+    pub fn new(provider: SharedProvider) -> Self {
        Self {
            extensions: Mutex::new(HashMap::new()),
            context: Mutex::new(PlatformExtensionContext {
@@ -251,9 +256,15 @@ impl ExtensionManager {
                extension_manager: None,
                tool_route_manager: None,
            }),
+            provider,
        }
    }

+    /// Create a new ExtensionManager with no provider (useful for tests)
+    pub fn new_without_provider() -> Self {
+        Self::new(Arc::new(Mutex::new(None)))
+    }
+
    pub async fn set_context(&self, context: PlatformExtensionContext) {
        *self.context.lock().await = context;
    }
@@ -348,6 +359,7 @@ impl ExtensionManager {
                        Duration::from_secs(
                            timeout.unwrap_or(crate::config::DEFAULT_EXTENSION_TIMEOUT),
                        ),
+                        self.provider.clone(),
                    )
                    .await?,
                )
@@ -388,6 +400,7 @@ impl ExtensionManager {
                    Duration::from_secs(
                        timeout.unwrap_or(crate::config::DEFAULT_EXTENSION_TIMEOUT),
                    ),
+                    self.provider.clone(),
                )
                .await;
                let client = if let Some(_auth_error) = extract_auth_error(&client_res) {
@@ -407,6 +420,7 @@ impl ExtensionManager {
                        Duration::from_secs(
                            timeout.unwrap_or(crate::config::DEFAULT_EXTENSION_TIMEOUT),
                        ),
+                        self.provider.clone(),
                    )
                    .await?
                } else {
@@ -430,7 +444,7 @@ impl ExtensionManager {
                // Check for malicious packages before launching the process
                extension_malware_check::deny_if_malicious_cmd_args(cmd, args).await?;

-                let client = child_process_client(command, timeout).await?;
+                let client = child_process_client(command, timeout, self.provider.clone()).await?;
                Box::new(client)
            }
            ExtensionConfig::Builtin {
@@ -459,7 +473,7 @@ impl ExtensionManager {
                let command = Command::new(cmd).configure(|command| {
                    command.arg("mcp").arg(name);
                });
-                let client = child_process_client(command, timeout).await?;
+                let client = child_process_client(command, timeout, self.provider.clone()).await?;
                Box::new(client)
            }
            ExtensionConfig::Platform { name, .. } => {
@@ -495,7 +509,7 @@ impl ExtensionManager {
                    command.arg("python").arg(file_path.to_str().unwrap());
                });

-                let client = child_process_client(command, timeout).await?;
+                let client = child_process_client(command, timeout, self.provider.clone()).await?;

                Box::new(client)
            }
@@ -1252,7 +1266,7 @@ mod tests {

    #[tokio::test]
    async fn test_get_client_for_tool() {
-        let extension_manager = ExtensionManager::new();
+        let extension_manager = ExtensionManager::new_without_provider();

        // Add some mock clients using the helper method
        extension_manager
@@ -1312,7 +1326,7 @@ mod tests {
    async fn test_dispatch_tool_call() {
        // test that dispatch_tool_call parses out the sanitized name correctly, and extracts
        // tool_names
-        let extension_manager = ExtensionManager::new();
+        let extension_manager = ExtensionManager::new_without_provider();

        // Add some mock clients using the helper method
        extension_manager
@@ -1429,7 +1443,7 @@ mod tests {

    #[tokio::test]
    async fn test_tool_availability_filtering() {
-        let extension_manager = ExtensionManager::new();
+        let extension_manager = ExtensionManager::new_without_provider();

        // Only "available_tool" should be available to the LLM
        let available_tools = vec!["available_tool".to_string()];
@@ -1457,7 +1471,7 @@ mod tests {

    #[tokio::test]
    async fn test_tool_availability_defaults_to_available() {
-        let extension_manager = ExtensionManager::new();
+        let extension_manager = ExtensionManager::new_without_provider();

        extension_manager
            .add_mock_extension_with_tools(
@@ -1482,7 +1496,7 @@ mod tests {

    #[tokio::test]
    async fn test_dispatch_unavailable_tool_returns_error() {
-        let extension_manager = ExtensionManager::new();
+        let extension_manager = ExtensionManager::new_without_provider();

        let available_tools = vec!["available_tool".to_string()];

@@ -1,21 +1,24 @@
-use rmcp::model::JsonObject;
+use crate::agents::types::SharedProvider;
+use rmcp::model::{Content, ErrorCode, JsonObject};
 /// MCP client implementation for Goose
 use rmcp::{
    model::{
        CallToolRequest, CallToolRequestParam, CallToolResult, CancelledNotification,
        CancelledNotificationMethod, CancelledNotificationParam, ClientCapabilities, ClientInfo,
-        ClientRequest, GetPromptRequest, GetPromptRequestParam, GetPromptResult, Implementation,
-        InitializeResult, ListPromptsRequest, ListPromptsResult, ListResourcesRequest,
-        ListResourcesResult, ListToolsRequest, ListToolsResult, LoggingMessageNotification,
+        ClientRequest, CreateMessageRequestParam, CreateMessageResult, GetPromptRequest,
+        GetPromptRequestParam, GetPromptResult, Implementation, InitializeResult,
+        ListPromptsRequest, ListPromptsResult, ListResourcesRequest, ListResourcesResult,
+        ListToolsRequest, ListToolsResult, LoggingMessageNotification,
        LoggingMessageNotificationMethod, PaginatedRequestParam, ProgressNotification,
        ProgressNotificationMethod, ProtocolVersion, ReadResourceRequest, ReadResourceRequestParam,
-        ReadResourceResult, RequestId, ServerNotification, ServerResult,
+        ReadResourceResult, RequestId, Role, SamplingMessage, ServerNotification, ServerResult,
    },
    service::{
-        ClientInitializeError, PeerRequestOptions, RequestHandle, RunningService, ServiceRole,
+        ClientInitializeError, PeerRequestOptions, RequestContext, RequestHandle, RunningService,
+        ServiceRole,
    },
    transport::IntoTransport,
-    ClientHandler, Peer, RoleClient, ServiceError, ServiceExt,
+    ClientHandler, ErrorData, Peer, RoleClient, ServiceError, ServiceExt,
 };
 use serde_json::Value;
 use std::{sync::Arc, time::Duration};
@@ -76,12 +79,17 @@ pub trait McpClientTrait: Send + Sync {

 pub struct GooseClient {
    notification_handlers: Arc<Mutex<Vec<Sender<ServerNotification>>>>,
+    provider: SharedProvider,
 }

 impl GooseClient {
-    pub fn new(handlers: Arc<Mutex<Vec<Sender<ServerNotification>>>>) -> Self {
+    pub fn new(
+        handlers: Arc<Mutex<Vec<Sender<ServerNotification>>>>,
+        provider: SharedProvider,
+    ) -> Self {
        GooseClient {
            notification_handlers: handlers,
+            provider,
        }
    }
 }
@@ -127,10 +135,88 @@ impl ClientHandler for GooseClient {
            });
    }

+    async fn create_message(
+        &self,
+        params: CreateMessageRequestParam,
+        _context: RequestContext<RoleClient>,
+    ) -> Result<CreateMessageResult, ErrorData> {
+        let provider = self
+            .provider
+            .lock()
+            .await
+            .as_ref()
+            .ok_or(ErrorData::new(
+                ErrorCode::INTERNAL_ERROR,
+                "Could not use provider",
+                None,
+            ))?
+            .clone();
+
+        let provider_ready_messages: Vec<crate::conversation::message::Message> = params
+            .messages
+            .iter()
+            .map(|msg| {
+                let base = match msg.role {
+                    Role::User => crate::conversation::message::Message::user(),
+                    Role::Assistant => crate::conversation::message::Message::assistant(),
+                };
+
+                match msg.content.as_text() {
+                    Some(text) => base.with_text(&text.text),
+                    None => base.with_content(msg.content.clone().into()),
+                }
+            })
+            .collect();
+
+        let system_prompt = params
+            .system_prompt
+            .as_deref()
+            .unwrap_or("You are a general-purpose AI agent called goose");
+
+        let (response, usage) = provider
+            .complete(system_prompt, &provider_ready_messages, &[])
+            .await
+            .map_err(|e| {
+                ErrorData::new(
+                    ErrorCode::INTERNAL_ERROR,
+                    "Unexpected error while completing the prompt",
+                    Some(Value::from(e.to_string())),
+                )
+            })?;
+
+        Ok(CreateMessageResult {
+            model: usage.model,
+            stop_reason: Some(CreateMessageResult::STOP_REASON_END_TURN.to_string()),
+            message: SamplingMessage {
+                role: Role::Assistant,
+                // TODO(alexhancock): MCP sampling currently only supports one content on each SamplingMessage
+                // https://modelcontextprotocol.io/specification/draft/client/sampling#messages
+                // This doesn't mesh well with goose's approach which has Vec<MessageContent>
+                // There is a proposal to MCP which is agreed to go in the next version to have SamplingMessages support multiple content parts
+                // https://github.com/modelcontextprotocol/modelcontextprotocol/pull/198
+                // Until that is formalized, we can take the first message content from the provider and use it
+                content: if let Some(content) = response.content.first() {
+                    match content {
+                        crate::conversation::message::MessageContent::Text(text) => {
+                            Content::text(&text.text)
+                        }
+                        crate::conversation::message::MessageContent::Image(img) => {
+                            Content::image(&img.data, &img.mime_type)
+                        }
+                        // TODO(alexhancock) - Content::Audio? goose's messages don't currently have it
+                        _ => Content::text(""),
+                    }
+                } else {
+                    Content::text("")
+                },
+            },
+        })
+    }
+
    fn get_info(&self) -> ClientInfo {
        ClientInfo {
            protocol_version: ProtocolVersion::V_2025_03_26,
-            capabilities: ClientCapabilities::builder().build(),
+            capabilities: ClientCapabilities::builder().enable_sampling().build(),
            client_info: Implementation {
                name: "goose".to_string(),
                version: std::env::var("GOOSE_MCP_CLIENT_VERSION")
@@ -155,6 +241,7 @@ impl McpClient {
    pub async fn connect<T, E, A>(
        transport: T,
        timeout: std::time::Duration,
+        provider: SharedProvider,
    ) -> Result<Self, ClientInitializeError>
    where
        T: IntoTransport<RoleClient, E, A>,
@@ -163,7 +250,7 @@ impl McpClient {
        let notification_subscribers =
            Arc::new(Mutex::new(Vec::<mpsc::Sender<ServerNotification>>::new()));

-        let client = GooseClient::new(notification_subscribers.clone());
+        let client = GooseClient::new(notification_subscribers.clone(), provider);
        let client: rmcp::service::RunningService<rmcp::RoleClient, GooseClient> =
            client.serve(transport).await?;
        let server_info = client.peer_info().cloned();
@@ -1,4 +1,5 @@
 use crate::mcp_utils::ToolResult;
+use crate::providers::base::Provider;
 use rmcp::model::{Content, Tool};
 use serde::{Deserialize, Serialize};
 use std::path::PathBuf;
@@ -9,6 +10,13 @@ use utoipa::ToSchema;
 /// Type alias for the tool result channel receiver
 pub type ToolResultReceiver = Arc<Mutex<mpsc::Receiver<(String, ToolResult<Vec<Content>>)>>>;

+/// This is used when we want to share the agent's current provider
+/// There are a lot of components to this definition so breaking it down:
+/// `Arc` enables shared ownership across threads or async contexts
+/// `Mutex` ensures mutable access is synchronized
+/// `Option` represents that a provider may or may not be set
+pub type SharedProvider = Arc<Mutex<Option<Arc<dyn Provider>>>>;
+
 /// Default timeout for retry operations (5 minutes)
 pub const DEFAULT_RETRY_TIMEOUT_SECONDS: u64 = 300;

@@ -2,6 +2,7 @@ use serde::Deserialize;
 use std::collections::HashMap;
 use std::fs::File;
 use std::path::PathBuf;
+use std::sync::Arc;
 use std::{env, fs};

 use rmcp::model::{CallToolRequestParam, Content};
@@ -10,6 +11,8 @@ use tokio_util::sync::CancellationToken;

 use goose::agents::extension::{Envs, ExtensionConfig};
 use goose::agents::extension_manager::ExtensionManager;
+use goose::model::ModelConfig;
+use goose::providers::openai::OpenAiProvider;

 use test_case::test_case;

@@ -72,6 +75,17 @@ enum TestMode {
    Playback,
 }

+// Use an export OPENAI_API_KEY when recording
+async fn create_recording_provider() -> Result<
+    Arc<tokio::sync::Mutex<Option<Arc<dyn goose::providers::base::Provider>>>>,
+    Box<dyn std::error::Error>,
+> {
+    let provider = OpenAiProvider::from_env(ModelConfig::new("gpt-5-mini")?).await?;
+    Ok(Arc::new(tokio::sync::Mutex::new(Some(
+        Arc::new(provider) as Arc<dyn goose::providers::base::Provider>
+    ))))
+}
+
 #[test_case(
    vec!["npx", "-y", "@modelcontextprotocol/server-everything"],
    vec![
@@ -79,6 +93,7 @@ enum TestMode {
        CallToolRequestParam { name: "add".into(), arguments: Some(object!({"a": 1, "b": 2 })) },
        CallToolRequestParam { name: "longRunningOperation".into(), arguments: Some(object!({"duration": 1, "steps": 5 })) },
        CallToolRequestParam { name: "structuredContent".into(), arguments: Some(object!({"location": "11238"})) },
+        CallToolRequestParam { name: "sampleLLM".into(), arguments: Some(object!({"prompt": "Please provide a quote from The Great Gatsby", "maxTokens": 100 })) },
    ],
    vec![]
 )]
@@ -206,7 +221,19 @@ async fn test_replayed_session(
        available_tools: vec![],
    };

-    let extension_manager = ExtensionManager::new();
+    let extension_manager = if matches!(mode, TestMode::Record) {
+        match create_recording_provider().await {
+            Ok(provider) => ExtensionManager::new(provider),
+            Err(e) => {
+                eprintln!("Failed to create OpenAI provider: {:?}", e);
+                eprintln!("Skipping test - ensure OPENAI_API_KEY is configured");
+                return;
+            }
+        }
+    } else {
+        // In playback mode, we don't need a real provider
+        ExtensionManager::new_without_provider()
+    };

    #[allow(clippy::redundant_closure_call)]
    let result = (async || -> Result<(), Box<dyn std::error::Error>> {
@@ -90,7 +90,7 @@
  [
    {
      "type": "text",
-      "text": "Available windows:\nMenubar",
+      "text": "Available windows:\n\nItem-0\nItem-0\nItem-0\nItem-0\nItem-0\nItem-0\nItem-0\nItem-0\nItem-0\nItem-0\nItem-0\nBattery\nWiFi\nItem-0\nBentoBox\nSiri\nClock\nMenubar\nDock\n~/Development/goose\ngoose – mcp_integration_test.rs\nwhat is the fast version of gpt5? - Google Search\n\nChatGPT\nDesktop\ntests\nDesktop\n+1 (310) 869-7623\n* cmoulton-office (Channel) - Block, Inc. - 1 new item - Slack\nNotes\n#🌌┃ecosystem | goose - Discord\nLock Screen — 1Password",
      "annotations": {
        "audience": [
          "assistant"
@@ -99,7 +99,7 @@
    },
    {
      "type": "text",
-      "text": "Available windows:\nMenubar",
+      "text": "Available windows:\n\nItem-0\nItem-0\nItem-0\nItem-0\nItem-0\nItem-0\nItem-0\nItem-0\nItem-0\nItem-0\nItem-0\nBattery\nWiFi\nItem-0\nBentoBox\nSiri\nClock\nMenubar\nDock\n~/Development/goose\ngoose – mcp_integration_test.rs\nwhat is the fast version of gpt5? - Google Search\n\nChatGPT\nDesktop\ntests\nDesktop\n+1 (310) 869-7623\n* cmoulton-office (Channel) - Block, Inc. - 1 new item - Slack\nNotes\n#🌌┃ecosystem | goose - Discord\nLock Screen — 1Password",
      "annotations": {
        "audience": [
          "user"
@@ -1,29 +1,10 @@
-STDIN: {"jsonrpc":"2.0","id":0,"method":"initialize","params":{"protocolVersion":"2025-03-26","capabilities":{},"clientInfo":{"name":"goose","version":"0.0.0"}}}
-STDERR: 2025-09-26 23:13:04 - Starting npx setup script.
-STDERR: 2025-09-26 23:13:04 - Creating directory ~/.config/goose/mcp-hermit/bin if it does not exist.
-STDERR: 2025-09-26 23:13:04 - Changing to directory ~/.config/goose/mcp-hermit.
-STDERR: 2025-09-26 23:13:04 - Hermit binary already exists. Skipping download.
-STDERR: 2025-09-26 23:13:04 - setting hermit cache to be local for MCP servers
-STDERR: 2025-09-26 23:13:04 - Updated PATH to include ~/.config/goose/mcp-hermit/bin.
-STDERR: 2025-09-26 23:13:04 - Checking for hermit in PATH.
-STDERR: 2025-09-26 23:13:04 - Initializing hermit.
-STDERR: 2025-09-26 23:13:04 - Installing Node.js with hermit.
-STDERR: 2025-09-26 23:13:04 - Verifying installation locations:
-STDERR: 2025-09-26 23:13:04 - hermit: /Users/angiej/.config/goose/mcp-hermit/bin/hermit
-STDERR: 2025-09-26 23:13:04 - node: /Users/angiej/.config/goose/mcp-hermit/bin/node
-STDERR: 2025-09-26 23:13:04 - npx: /Users/angiej/.config/goose/mcp-hermit/bin/npx
-STDERR: 2025-09-26 23:13:04 - Checking for GOOSE_NPM_REGISTRY and GOOSE_NPM_CERT environment variables for custom npm registry setup...
-STDERR: 2025-09-26 23:13:05 - Checking custom goose registry availability: https://global.block-artifacts.com/artifactory/api/npm/square-npm/
-STDERR: 2025-09-26 23:13:05 - https://global.block-artifacts.com/artifactory/api/npm/square-npm/ is accessible. Using it for npm registry.
-STDERR: 2025-09-26 23:13:06 - Downloading certificate from: https://block-ca.squareup.com/root-certs.pem
-STDERR: 2025-09-26 23:13:07 - Certificate downloaded successfully.
-STDERR: 2025-09-26 23:13:07 - Executing 'npx' command with arguments: -y @modelcontextprotocol/server-everything
+STDIN: {"jsonrpc":"2.0","id":0,"method":"initialize","params":{"protocolVersion":"2025-03-26","capabilities":{"sampling":{}},"clientInfo":{"name":"goose","version":"0.0.0"}}}
 STDERR: Starting default (STDIO) server...
 STDERR: Starting logs update interval
 STDOUT: {"result":{"protocolVersion":"2025-03-26","capabilities":{"prompts":{},"resources":{"subscribe":true},"tools":{},"logging":{},"completions":{}},"serverInfo":{"name":"example-servers/everything","title":"Everything Example Server","version":"1.0.0"},"instructions":"Testing and demonstration server for MCP protocol features.\n\n## Resources\n\nResources 1-100 follow pattern: even IDs contain text, odd IDs contain binary data. Resources are paginated at 10 items per page with cursor-based navigation.\n\n## Key dependencies\n\nProgress notifications require `_meta.progressToken` in tool calls. Resource subscriptions generate updates every 10 seconds.\n\n## Performance characteristics\n\nServer generates automatic log messages every 20 seconds (filtered by current log level), stderr notifications every 30 seconds, and resource update notifications every 10 seconds for active subscriptions. Log level changes affect message filtering in real-time.\n\n## Multi-modal testing\n\n`complex_prompt` includes both text arguments and image content for testing client multi-modal handling. `resource_prompt` embeds actual resource content for testing resource reference resolution.\n\nArgument completion is available for prompt parameters and resource IDs. Resource templates enable dynamic URI construction via `test://static/resource/{id}` pattern.\n\n## Easter egg\n\nIf asked about server instructions, respond with \"🎉 Server instructions are working! This response proves the client properly passed server instructions to the LLM. This demonstrates MCP's instructions feature in action.\"\n"},"jsonrpc":"2.0","id":0}
 STDIN: {"jsonrpc":"2.0","method":"notifications/initialized"}
-STDIN: {"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"_meta":{"progressToken":0},"name":"echo","arguments":{"message":"Hello, world!"}}}
 STDOUT: {"method":"notifications/message","params":{"level":"info","logger":"everything-server","data":"Client does not support MCP roots protocol"},"jsonrpc":"2.0"}
+STDIN: {"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"_meta":{"progressToken":0},"name":"echo","arguments":{"message":"Hello, world!"}}}
 STDOUT: {"result":{"content":[{"type":"text","text":"Echo: Hello, world!"}]},"jsonrpc":"2.0","id":1}
 STDIN: {"jsonrpc":"2.0","id":2,"method":"tools/call","params":{"_meta":{"progressToken":1},"name":"add","arguments":{"a":1,"b":2}}}
 STDOUT: {"result":{"content":[{"type":"text","text":"The sum of 1 and 2 is 3."}]},"jsonrpc":"2.0","id":2}
@@ -36,5 +17,9 @@ STDOUT: {"method":"notifications/progress","params":{"progress":5,"total":5,"pro
 STDOUT: {"result":{"content":[{"type":"text","text":"Long running operation completed. Duration: 1 seconds, Steps: 5."}]},"jsonrpc":"2.0","id":3}
 STDIN: {"jsonrpc":"2.0","id":4,"method":"tools/call","params":{"_meta":{"progressToken":3},"name":"structuredContent","arguments":{"location":"11238"}}}
 STDOUT: {"result":{"content":[{"type":"text","text":"{\"temperature\":22.5,\"conditions\":\"Partly cloudy\",\"humidity\":65}"}],"structuredContent":{"temperature":22.5,"conditions":"Partly cloudy","humidity":65}},"jsonrpc":"2.0","id":4}
-STDOUT: {"method":"notifications/message","params":{"level":"emergency","data":"Emergency-level message"},"jsonrpc":"2.0"}
-STDERR: node:events:497
+STDIN: {"jsonrpc":"2.0","id":5,"method":"tools/call","params":{"_meta":{"progressToken":4},"name":"sampleLLM","arguments":{"maxTokens":100,"prompt":"Please provide a quote from The Great Gatsby"}}}
+STDOUT: {"method":"sampling/createMessage","params":{"messages":[{"role":"user","content":{"type":"text","text":"Resource sampleLLM context: Please provide a quote from The Great Gatsby"}}],"systemPrompt":"You are a helpful test server.","maxTokens":100,"temperature":0.7,"includeContext":"thisServer"},"jsonrpc":"2.0","id":0}
+STDIN: {"jsonrpc":"2.0","id":0,"result":{"model":"gpt-5-mini-2025-08-07","stopReason":"endTurn","role":"assistant","content":{"type":"text","text":"\"So we beat on, boats against the current, borne back ceaselessly into the past.\" — F. Scott Fitzgerald, The Great Gatsby."}}}
+STDOUT: {"result":{"content":[{"type":"text","text":"LLM sampling result: \"So we beat on, boats against the current, borne back ceaselessly into the past.\" — F. Scott Fitzgerald, The Great Gatsby."}]},"jsonrpc":"2.0","id":5}
+STDOUT: {"method":"notifications/message","params":{"level":"error","data":"Error-level message"},"jsonrpc":"2.0"}
+STDERR: node:events:486
@@ -22,5 +22,11 @@
      "type": "text",
      "text": "{\"temperature\":22.5,\"conditions\":\"Partly cloudy\",\"humidity\":65}"
    }
+  ],
+  [
+    {
+      "type": "text",
+      "text": "LLM sampling result: \"So we beat on, boats against the current, borne back ceaselessly into the past.\" — F. Scott Fitzgerald, The Great Gatsby."
+    }
  ]
 ]
@@ -1,29 +1,5 @@
-STDIN: {"jsonrpc":"2.0","id":0,"method":"initialize","params":{"protocolVersion":"2025-03-26","capabilities":{},"clientInfo":{"name":"goose","version":"0.0.0"}}}
-STDERR: 2025-09-26 23:13:04 - Starting uvx setup script.
-STDERR: 2025-09-26 23:13:04 - Creating directory ~/.config/goose/mcp-hermit/bin if it does not exist.
-STDERR: 2025-09-26 23:13:04 - Changing to directory ~/.config/goose/mcp-hermit.
-STDERR: 2025-09-26 23:13:04 - Hermit binary already exists. Skipping download.
-STDERR: 2025-09-26 23:13:04 - setting hermit cache to be local for MCP servers
-STDERR: 2025-09-26 23:13:04 - Updated PATH to include ~/.config/goose/mcp-hermit/bin.
-STDERR: 2025-09-26 23:13:04 - Checking for hermit in PATH.
-STDERR: 2025-09-26 23:13:04 - Initializing hermit.
-STDERR: 2025-09-26 23:13:04 - hermit install python 3.10
-STDERR: 2025-09-26 23:13:04 - Installing UV with hermit.
-STDERR: 2025-09-26 23:13:04 - Verifying installation locations:
-STDERR: 2025-09-26 23:13:04 - hermit: /Users/angiej/.config/goose/mcp-hermit/bin/hermit
-STDERR: 2025-09-26 23:13:04 - uv: /Users/angiej/.config/goose/mcp-hermit/bin/uv
-STDERR: 2025-09-26 23:13:04 - uvx: /Users/angiej/.config/goose/mcp-hermit/bin/uvx
-STDERR: 2025-09-26 23:13:04 - Checking for GOOSE_UV_REGISTRY environment variable for custom python/pip/UV registry setup...
-STDERR: 2025-09-26 23:13:05 - Checking custom goose registry availability: https://global.block-artifacts.com/artifactory/api/pypi/block-pypi/simple
-STDERR: 2025-09-26 23:13:05 - https://global.block-artifacts.com/artifactory/api/pypi/block-pypi/simple is accessible, setting it as UV_DEFAULT_INDEX. Setting UV_NATIVE_TLS to true.
-STDERR: 2025-09-26 23:13:05 - Executing 'uvx' command with arguments: mcp-server-fetch
-STDOUT: {"jsonrpc":"2.0","id":0,"result":{"protocolVersion":"2025-03-26","capabilities":{"experimental":{},"prompts":{"listChanged":false},"tools":{"listChanged":false}},"serverInfo":{"name":"mcp-fetch","version":"1.15.0"}}}
+STDIN: {"jsonrpc":"2.0","id":0,"method":"initialize","params":{"protocolVersion":"2025-03-26","capabilities":{"sampling":{}},"clientInfo":{"name":"goose","version":"0.0.0"}}}
+STDOUT: {"jsonrpc":"2.0","id":0,"result":{"protocolVersion":"2025-03-26","capabilities":{"experimental":{},"prompts":{"listChanged":false},"tools":{"listChanged":false}},"serverInfo":{"name":"mcp-fetch","version":"1.19.0"}}}
 STDIN: {"jsonrpc":"2.0","method":"notifications/initialized"}
 STDIN: {"jsonrpc":"2.0","id":1,"method":"tools/call","params":{"_meta":{"progressToken":0},"name":"fetch","arguments":{"url":"https://example.com"}}}
-STDERR: npm error code FETCH_ERROR
-STDERR: npm error errno FETCH_ERROR
-STDERR: npm error invalid json response body at https://blocked.teams.cloudflare.com/?account_id=1e25787f854fa4b713d08a859d3e16ed&background_color=%23000000&block_reason=This+has+been+blocked+as+part+of+the+Dependency+Confusion+threat.+Please+see+go%2Fdependencyconfusionpypi+and+go%2Fdependencyconfusionnpm+for+more+info.&device_id=***&footer_text=The+website+you+are+trying+to+access+has+been+blocked+because+it+presents+a+risk+to+the+safety+and+security+of+Block%E2%80%99s+IT+systems.&header_text=This+page+presents+a+risk+to+Block&location=cf1ebd1203624140846ced63a200519e&logo_path=https%3A%2F%2Fmedia.block.xyz%2Flogos%2Fblock-jewel_white.png&mailto_address=&mailto_subject=&name=Block%2C+Inc.&params_sign=yrMcT5HYDMHvixy%2BdLHApce3BcNYIdlI8qh3wTcIrLA%3D&query_id=***&rule_id=***&source_ip=2a09%3Abac0%3A1000%3A2df%3A%3A281%3Ac0&suppress_footer=false&url=registry.npmjs.org&user_id=*** reason: Unexpected token '<', "
-STDERR: npm error <!DOCTYPE "... is not valid JSON
-STDERR: npm error A complete log of this run can be found in: /Users/angiej/.config/goose/mcp-hermit/.hermit/node/cache/_logs/2025-09-27T04_13_13_364Z-debug-0.log
-STDOUT: {"jsonrpc":"2.0","id":1,"result":{"content":[{"type":"text","text":"Command '['npm', 'install']' returned non-zero exit status 1."}],"isError":true}}
-STDERR: 2025-09-26 23:13:14 - uvx setup script completed successfully.
+STDOUT: {"jsonrpc":"2.0","id":1,"result":{"content":[{"type":"text","text":"Contents of https://example.com/:\nThis domain is for use in documentation examples without needing permission. Avoid use in operations.\n\n[Learn more](https://iana.org/domains/example)"}],"isError":false}}
@@ -2,7 +2,7 @@
  [
    {
      "type": "text",
-      "text": "Command '['npm', 'install']' returned non-zero exit status 1."
+      "text": "Contents of https://example.com/:\nThis domain is for use in documentation examples without needing permission. Avoid use in operations.\n\n[Learn more](https://iana.org/domains/example)"
    }
  ]
 ]
@@ -0,0 +1,81 @@
+#!/bin/bash
+if [ -f .env ]; then
+  export $(grep -v '^#' .env | xargs)
+fi
+
+if [ -z "$SKIP_BUILD" ]; then
+  echo "Building goose..."
+  cargo build --release --bin goose
+  echo ""
+else
+  echo "Skipping build (SKIP_BUILD is set)..."
+  echo ""
+fi
+
+SCRIPT_DIR=$(pwd)
+
+PROVIDERS=(
+  "openrouter:google/gemini-2.5-pro:google/gemini-2.5-flash:anthropic/claude-sonnet-4.5:qwen/qwen3-coder"
+  "openai:gpt-4o:gpt-4o-mini:gpt-3.5-turbo:gpt-5"
+  "anthropic:claude-sonnet-4-5-20250929:claude-opus-4-1-20250805"
+  "google:gemini-2.5-pro:gemini-2.5-flash"
+  "tetrate:claude-sonnet-4-20250514"
+)
+
+# In CI, only run Databricks tests if DATABRICKS_HOST and DATABRICKS_TOKEN are set
+# Locally, always run Databricks tests
+if [ -n "$CI" ]; then
+  if [ -n "$DATABRICKS_HOST" ] && [ -n "$DATABRICKS_TOKEN" ]; then
+    echo "✓ Including Databricks tests"
+    PROVIDERS+=("databricks:databricks-claude-sonnet-4:gemini-2-5-flash:gpt-4o")
+  else
+    echo "⚠️  Skipping Databricks tests (DATABRICKS_HOST and DATABRICKS_TOKEN required in CI)"
+  fi
+else
+  echo "✓ Including Databricks tests"
+  PROVIDERS+=("databricks:databricks-claude-sonnet-4:gemini-2-5-flash:gpt-4o")
+fi
+
+RESULTS=()
+
+for provider_config in "${PROVIDERS[@]}"; do
+  IFS=':' read -ra PARTS <<< "$provider_config"
+  PROVIDER="${PARTS[0]}"
+  for i in $(seq 1 $((${#PARTS[@]} - 1))); do
+    MODEL="${PARTS[$i]}"
+    export GOOSE_PROVIDER="$PROVIDER"
+    export GOOSE_MODEL="$MODEL"
+    TESTDIR=$(mktemp -d)
+    echo "Provider: ${PROVIDER}"
+    echo "Model: ${MODEL}"
+    echo ""
+    TMPFILE=$(mktemp)
+    (cd "$TESTDIR" && "$SCRIPT_DIR/target/release/goose" run --text "Use the everything__sampleLLM tool to ask for a quote from The Great Gatsby" --with-extension "npx -y @modelcontextprotocol/server-everything" 2>&1) | tee "$TMPFILE"
+    echo ""
+    if grep -q "sampleLLM | everything" "$TMPFILE"; then
+      echo "✓ SUCCESS: MCP sampling test passed - sampleLLM tool called"
+      RESULTS+=("✓ MCP Sampling ${PROVIDER}: ${MODEL}")
+    else
+      echo "✗ FAILED: MCP sampling test failed - sampleLLM tool not called"
+      RESULTS+=("✗ MCP Sampling ${PROVIDER}: ${MODEL}")
+    fi
+    rm "$TMPFILE"
+    rm -rf "$TESTDIR"
+    echo "---"
+  done
+done
+
+echo ""
+echo "=== MCP Sampling Test Summary ==="
+for result in "${RESULTS[@]}"; do
+  echo "$result"
+done
+
+if echo "${RESULTS[@]}" | grep -q "✗"; then
+  echo ""
+  echo "Some MCP sampling tests failed!"
+  exit 1
+else
+  echo ""
+  echo "All MCP sampling tests passed!"
+fi