test(mcp): add image tool test and consolidate MCP test fixtures (#7019)

Signed-off-by: Adrian Cole <adrian@tetrate.io>

crates/goose-test-support/src/lib.rs

@@ -0,0 +1,5 @@
+pub mod mcp;
+pub mod session;
+
+pub use mcp::{McpFixture, FAKE_CODE, TEST_IMAGE_B64};
+pub use session::{ExpectedSessionId, TEST_SESSION_ID};

crates/goose-test-support/src/mcp.rs

@@ -0,0 +1,180 @@
+use crate::session::{ExpectedSessionId, SESSION_ID_HEADER};
+use rmcp::model::{
+    CallToolResult, ClientNotification, ClientRequest, Content, ErrorCode, Implementation, Meta,
+    ProtocolVersion, ServerCapabilities, ServerInfo,
+};
+use rmcp::service::{DynService, NotificationContext, RequestContext, ServiceExt, ServiceRole};
+use rmcp::transport::streamable_http_server::{
+    session::local::LocalSessionManager, StreamableHttpServerConfig, StreamableHttpService,
+};
+use rmcp::{
+    handler::server::router::tool::ToolRouter, tool, tool_handler, tool_router,
+    ErrorData as McpError, RoleServer, ServerHandler, Service,
+};
+use tokio::task::JoinHandle;
+
+pub const FAKE_CODE: &str = "test-uuid-12345-67890";
+
+pub const TEST_IMAGE_B64: &str = include_str!("test_assets/test_image.b64").trim_ascii_end();
+
+pub trait HasMeta {
+    fn meta(&self) -> &Meta;
+}
+
+impl<R: ServiceRole> HasMeta for RequestContext<R> {
+    fn meta(&self) -> &Meta {
+        &self.meta
+    }
+}
+
+impl<R: ServiceRole> HasMeta for NotificationContext<R> {
+    fn meta(&self) -> &Meta {
+        &self.meta
+    }
+}
+
+struct ValidatingService<S> {
+    inner: S,
+    expected_session_id: ExpectedSessionId,
+}
+
+impl<S> ValidatingService<S> {
+    fn new(inner: S, expected_session_id: ExpectedSessionId) -> Self {
+        Self {
+            inner,
+            expected_session_id,
+        }
+    }
+
+    fn validate<C: HasMeta>(&self, context: &C) -> Result<(), McpError> {
+        let actual = context
+            .meta()
+            .0
+            .get(SESSION_ID_HEADER)
+            .and_then(|v| v.as_str());
+        self.expected_session_id
+            .validate(actual)
+            .map_err(|e| McpError::new(ErrorCode::INVALID_REQUEST, e, None))
+    }
+}
+
+impl<S: Service<RoleServer>> Service<RoleServer> for ValidatingService<S> {
+    async fn handle_request(
+        &self,
+        request: ClientRequest,
+        context: RequestContext<RoleServer>,
+    ) -> Result<rmcp::model::ServerResult, McpError> {
+        if !matches!(request, ClientRequest::InitializeRequest(_)) {
+            self.validate(&context)?;
+        }
+        self.inner.handle_request(request, context).await
+    }
+
+    async fn handle_notification(
+        &self,
+        notification: ClientNotification,
+        context: NotificationContext<RoleServer>,
+    ) -> Result<(), McpError> {
+        if !matches!(notification, ClientNotification::InitializedNotification(_)) {
+            self.validate(&context).ok();
+        }
+        self.inner.handle_notification(notification, context).await
+    }
+
+    fn get_info(&self) -> ServerInfo {
+        self.inner.get_info()
+    }
+}
+
+#[derive(Clone)]
+pub struct McpFixtureServer {
+    tool_router: ToolRouter<McpFixtureServer>,
+}
+
+impl Default for McpFixtureServer {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+#[tool_router]
+impl McpFixtureServer {
+    pub fn new() -> Self {
+        Self {
+            tool_router: Self::tool_router(),
+        }
+    }
+
+    #[tool(description = "Get the code")]
+    fn get_code(&self) -> Result<CallToolResult, McpError> {
+        Ok(CallToolResult::success(vec![Content::text(FAKE_CODE)]))
+    }
+
+    #[tool(description = "Get an image")]
+    fn get_image(&self) -> Result<CallToolResult, McpError> {
+        Ok(CallToolResult::success(vec![Content::image(
+            TEST_IMAGE_B64,
+            "image/png",
+        )]))
+    }
+}
+
+#[tool_handler]
+impl ServerHandler for McpFixtureServer {
+    fn get_info(&self) -> ServerInfo {
+        ServerInfo {
+            protocol_version: ProtocolVersion::V_2025_03_26,
+            capabilities: ServerCapabilities::builder().enable_tools().build(),
+            server_info: Implementation {
+                name: "mcp-fixture".into(),
+                version: "1.0.0".into(),
+                ..Default::default()
+            },
+            instructions: Some("Test server with get_code and get_image tools.".into()),
+        }
+    }
+}
+
+pub struct McpFixture {
+    pub url: String,
+    handle: JoinHandle<()>,
+}
+
+impl Drop for McpFixture {
+    fn drop(&mut self) {
+        self.handle.abort();
+    }
+}
+
+type McpServiceFactory =
+    Box<dyn Fn() -> Result<Box<dyn DynService<RoleServer>>, std::io::Error> + Send + Sync>;
+
+impl McpFixture {
+    pub async fn new(expected_session_id: Option<ExpectedSessionId>) -> Self {
+        let service_factory: McpServiceFactory = match expected_session_id {
+            Some(expected_session_id) => Box::new(move || {
+                Ok(
+                    ValidatingService::new(McpFixtureServer::new(), expected_session_id.clone())
+                        .into_dyn(),
+                )
+            }),
+            None => Box::new(|| Ok(McpFixtureServer::new().into_dyn())),
+        };
+
+        let service = StreamableHttpService::new(
+            service_factory,
+            LocalSessionManager::default().into(),
+            StreamableHttpServerConfig::default(),
+        );
+        let router = axum::Router::new().nest_service("/mcp", service);
+        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
+        let addr = listener.local_addr().unwrap();
+        let url = format!("http://{addr}/mcp");
+
+        let handle = tokio::spawn(async move {
+            axum::serve(listener, router).await.unwrap();
+        });
+
+        Self { url, handle }
+    }
+}

crates/goose-test-support/src/session.rs

@@ -0,0 +1,56 @@
+use std::sync::{Arc, Mutex};
+
+pub const TEST_SESSION_ID: &str = "test-session-id";
+
+const NOT_YET_SET: &str = "session-id-not-yet-set";
+pub(crate) const SESSION_ID_HEADER: &str = "agent-session-id";
+
+#[derive(Clone)]
+pub struct ExpectedSessionId {
+    value: Arc<Mutex<String>>,
+    errors: Arc<Mutex<Vec<String>>>,
+}
+
+impl Default for ExpectedSessionId {
+    fn default() -> Self {
+        Self {
+            value: Arc::new(Mutex::new(NOT_YET_SET.to_string())),
+            errors: Arc::new(Mutex::new(Vec::new())),
+        }
+    }
+}
+
+impl ExpectedSessionId {
+    pub fn set(&self, id: impl Into<String>) {
+        *self.value.lock().unwrap() = id.into();
+    }
+
+    pub fn validate(&self, actual: Option<&str>) -> Result<(), String> {
+        let expected = self.value.lock().unwrap();
+        let err = match actual {
+            Some(act) if act == *expected => None,
+            _ => Some(format!(
+                "{} mismatch: expected '{}', got {:?}",
+                SESSION_ID_HEADER, expected, actual
+            )),
+        };
+        match err {
+            Some(e) => {
+                self.errors.lock().unwrap().push(e.clone());
+                Err(e)
+            }
+            None => Ok(()),
+        }
+    }
+
+    pub fn assert_matches(&self, actual: &str) {
+        let result = self.validate(Some(actual));
+        assert!(result.is_ok(), "{}", result.unwrap_err());
+        let errors = self.errors.lock().unwrap();
+        assert!(
+            errors.is_empty(),
+            "Session ID validation errors: {:?}",
+            *errors
+        );
+    }
+}