feat(acp): add session/set_mode handler (#7801)

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

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

@@ -2,4 +2,6 @@ pub mod mcp;
 pub mod session;
 
 pub use mcp::{McpFixture, FAKE_CODE, TEST_IMAGE_B64};
-pub use session::{ExpectedSessionId, TEST_MODEL, TEST_SESSION_ID};
+pub use session::{
+    EnforceSessionId, ExpectedSessionId, IgnoreSessionId, TEST_MODEL, TEST_SESSION_ID,
+};

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

@@ -1,4 +1,5 @@
-use crate::session::{ExpectedSessionId, SESSION_ID_HEADER};
+use crate::session::SESSION_ID_HEADER;
+use crate::ExpectedSessionId;
 use rmcp::model::{
     CallToolResult, ClientNotification, ClientRequest, Content, ErrorCode, Implementation,
     InitializeResult, Meta, ProtocolVersion, ServerCapabilities, ServerInfo,
@@ -11,6 +12,7 @@ use rmcp::{
     handler::server::router::tool::ToolRouter, tool, tool_handler, tool_router,
     ErrorData as McpError, RoleServer, ServerHandler, Service,
 };
+use std::sync::Arc;
 use tokio::task::JoinHandle;
 
 pub const FAKE_CODE: &str = "test-uuid-12345-67890";
@@ -35,11 +37,11 @@ impl<R: ServiceRole> HasMeta for NotificationContext<R> {
 
 struct ValidatingService<S> {
     inner: S,
-    expected_session_id: ExpectedSessionId,
+    expected_session_id: Arc<dyn ExpectedSessionId>,
 }
 
 impl<S> ValidatingService<S> {
-    fn new(inner: S, expected_session_id: ExpectedSessionId) -> Self {
+    fn new(inner: S, expected_session_id: Arc<dyn ExpectedSessionId>) -> Self {
         Self {
             inner,
             expected_session_id,
@@ -144,16 +146,13 @@ 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())),
-        };
+    pub async fn new(expected_session_id: Arc<dyn ExpectedSessionId>) -> Self {
+        let service_factory: McpServiceFactory = Box::new(move || {
+            Ok(
+                ValidatingService::new(McpFixtureServer::new(), expected_session_id.clone())
+                    .into_dyn(),
+            )
+        });
 
         let service = StreamableHttpService::new(
             service_factory,

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

@@ -6,13 +6,19 @@ pub const TEST_MODEL: &str = "gpt-5-nano";
 const NOT_YET_SET: &str = "session-id-not-yet-set";
 pub(crate) const SESSION_ID_HEADER: &str = "agent-session-id";
 
+pub trait ExpectedSessionId: Send + Sync {
+    fn set(&self, id: &str);
+    fn validate(&self, actual: Option<&str>) -> Result<(), String>;
+    fn assert_matches(&self, actual: &str);
+}
+
 #[derive(Clone)]
-pub struct ExpectedSessionId {
+pub struct EnforceSessionId {
     value: Arc<Mutex<String>>,
     errors: Arc<Mutex<Vec<String>>>,
 }
 
-impl Default for ExpectedSessionId {
+impl Default for EnforceSessionId {
     fn default() -> Self {
         Self {
             value: Arc::new(Mutex::new(NOT_YET_SET.to_string())),
@@ -21,12 +27,12 @@ impl Default for ExpectedSessionId {
     }
 }
 
-impl ExpectedSessionId {
-    pub fn set(&self, id: impl Into<String>) {
+impl ExpectedSessionId for EnforceSessionId {
+    fn set(&self, id: &str) {
         *self.value.lock().unwrap() = id.into();
     }
 
-    pub fn validate(&self, actual: Option<&str>) -> Result<(), String> {
+    fn validate(&self, actual: Option<&str>) -> Result<(), String> {
         let expected = self.value.lock().unwrap();
         let err = match actual {
             Some(act) if act == *expected => None,
@@ -44,7 +50,7 @@ impl ExpectedSessionId {
         }
     }
 
-    pub fn assert_matches(&self, actual: &str) {
+    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();
@@ -55,3 +61,14 @@ impl ExpectedSessionId {
         );
     }
 }
+
+#[derive(Clone)]
+pub struct IgnoreSessionId;
+
+impl ExpectedSessionId for IgnoreSessionId {
+    fn set(&self, _id: &str) {}
+    fn validate(&self, _actual: Option<&str>) -> Result<(), String> {
+        Ok(())
+    }
+    fn assert_matches(&self, _actual: &str) {}
+}