diff --git a/Cargo.lock b/Cargo.lock index 348b71dd..e56eab1e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4347,6 +4347,7 @@ dependencies = [ "schemars 1.2.1", "serde", "serde_json", + "strum", "tempfile", "test-case", "tokio", diff --git a/crates/goose-acp/Cargo.toml b/crates/goose-acp/Cargo.toml index 41d63633..04aac357 100644 --- a/crates/goose-acp/Cargo.toml +++ b/crates/goose-acp/Cargo.toml @@ -37,6 +37,7 @@ serde_json = { workspace = true } futures = { workspace = true } regex = { workspace = true } fs-err = "3" +strum = { workspace = true } url = { workspace = true } # HTTP server dependencies diff --git a/crates/goose-acp/src/server.rs b/crates/goose-acp/src/server.rs index ef8413aa..ca5f475e 100644 --- a/crates/goose-acp/src/server.rs +++ b/crates/goose-acp/src/server.rs @@ -13,7 +13,7 @@ use goose::config::base::CONFIG_YAML_NAME; use goose::config::extensions::get_enabled_extensions_with_config; use goose::config::paths::Paths; use goose::config::permission::PermissionManager; -use goose::config::Config; +use goose::config::{Config, GooseMode}; use goose::conversation::message::{ActionRequiredData, Message, MessageContent}; use goose::conversation::Conversation; use goose::mcp_utils::ToolResult; @@ -27,16 +27,17 @@ use goose_acp_macros::custom_methods; use rmcp::model::{CallToolResult, RawContent, ResourceContents, Role}; use sacp::schema::{ AgentCapabilities, AuthMethod, AuthenticateRequest, AuthenticateResponse, BlobResourceContents, - CancelNotification, Content, ContentBlock, ContentChunk, EmbeddedResource, + CancelNotification, Content, ContentBlock, ContentChunk, CurrentModeUpdate, EmbeddedResource, EmbeddedResourceResource, FileSystemCapability, ImageContent, InitializeRequest, InitializeResponse, ListSessionsResponse, LoadSessionRequest, LoadSessionResponse, McpCapabilities, McpServer, ModelId, ModelInfo, NewSessionRequest, NewSessionResponse, PermissionOption, PermissionOptionKind, PromptCapabilities, PromptRequest, PromptResponse, RequestPermissionOutcome, RequestPermissionRequest, ResourceLink, SessionCapabilities, - SessionId, SessionInfo, SessionListCapabilities, SessionModelState, SessionNotification, - SessionUpdate, SetSessionModelRequest, SetSessionModelResponse, StopReason, TextContent, - TextResourceContents, ToolCall, ToolCallContent, ToolCallId, ToolCallLocation, ToolCallStatus, - ToolCallUpdate, ToolCallUpdateFields, ToolKind, + SessionId, SessionInfo, SessionListCapabilities, SessionMode, SessionModeId, SessionModeState, + SessionModelState, SessionNotification, SessionUpdate, SetSessionModeRequest, + SetSessionModeResponse, SetSessionModelRequest, SetSessionModelResponse, StopReason, + TextContent, TextResourceContents, ToolCall, ToolCallContent, ToolCallId, ToolCallLocation, + ToolCallStatus, ToolCallUpdate, ToolCallUpdateFields, ToolKind, }; use sacp::{AgentToClient, ByteStreams, Handled, JrConnectionCx, JrMessageHandler, MessageCx}; use std::collections::HashMap; @@ -65,7 +66,7 @@ pub struct GooseAcpAgent { config_dir: std::path::PathBuf, session_manager: Arc, permission_manager: Arc, - goose_mode: goose::config::GooseMode, + goose_mode: GooseMode, disable_session_naming: bool, } @@ -322,6 +323,24 @@ async fn build_model_state(provider: &dyn Provider, current_model: &str) -> Sess ) } +fn build_mode_state(current_mode: GooseMode) -> Result { + use strum::{EnumMessage, VariantNames}; + let mut available = Vec::with_capacity(GooseMode::VARIANTS.len()); + for &name in GooseMode::VARIANTS { + let goose_mode: GooseMode = name.parse().map_err(|_| { + sacp::Error::internal_error() // impossible but satisfy linters + .data(format!("Failed to parse GooseMode variant: {}", name)) + })?; + let mut mode = SessionMode::new(SessionModeId::new(name), name); + mode.description = goose_mode.get_message().map(Into::into); + available.push(mode); + } + Ok(SessionModeState::new( + SessionModeId::new(current_mode.to_string()), + available, + )) +} + impl GooseAcpAgent { pub fn permission_manager(&self) -> Arc { Arc::clone(&self.permission_manager) @@ -332,7 +351,7 @@ impl GooseAcpAgent { builtins: Vec, data_dir: std::path::PathBuf, config_dir: std::path::PathBuf, - goose_mode: goose::config::GooseMode, + goose_mode: GooseMode, disable_session_naming: bool, ) -> Result { let session_manager = Arc::new(SessionManager::new(data_dir)); @@ -800,7 +819,6 @@ impl GooseAcpAgent { .map_err(|e| { sacp::Error::internal_error().data(format!("Failed to create agent: {}", e)) })?; - let provider = self .init_provider(&agent, &goose_session) .await @@ -823,13 +841,17 @@ impl GooseAcpAgent { info!( session_id = %goose_session.id, session_type = "acp", + goose_mode = %self.goose_mode, "Session started" ); let model_state = build_model_state(&*provider, &provider.get_model_config().model_name).await; + let mode_state = build_mode_state(self.goose_mode)?; - Ok(NewSessionResponse::new(SessionId::new(goose_session.id)).models(model_state)) + Ok(NewSessionResponse::new(SessionId::new(goose_session.id)) + .models(model_state) + .modes(mode_state)) } async fn init_provider(&self, agent: &Agent, session: &Session) -> Result> { @@ -848,6 +870,21 @@ impl GooseAcpAgent { Ok(provider) } + async fn get_session_agent( + &self, + session_id: &str, + cancel_token: Option, + ) -> Result, sacp::Error> { + let mut sessions = self.sessions.lock().await; + let session = sessions.get_mut(session_id).ok_or_else(|| { + sacp::Error::invalid_params().data(format!("Session not found: {}", session_id)) + })?; + if let Some(token) = cancel_token { + session.cancel_token = Some(token); + } + Ok(session.agent.clone()) + } + async fn add_mcp_extensions( agent: &Agent, mcp_servers: Vec, @@ -976,16 +1013,23 @@ impl GooseAcpAgent { let mut sessions = self.sessions.lock().await; sessions.insert(session_id.clone(), session); + // TODO: read from goose_session.goose_mode after #7603 + let goose_mode = self.goose_mode; + info!( session_id = %session_id, session_type = "acp", + goose_mode = %goose_mode, "Session loaded" ); let model_state = build_model_state(&*provider, &provider.get_model_config().model_name).await; + let mode_state = build_mode_state(goose_mode)?; - Ok(LoadSessionResponse::new().models(model_state)) + Ok(LoadSessionResponse::new() + .models(model_state) + .modes(mode_state)) } async fn on_prompt( @@ -996,14 +1040,9 @@ impl GooseAcpAgent { let session_id = args.session_id.0.to_string(); let cancel_token = CancellationToken::new(); - let agent = { - let mut sessions = self.sessions.lock().await; - let session = sessions.get_mut(&session_id).ok_or_else(|| { - sacp::Error::invalid_params().data(format!("Session not found: {}", session_id)) - })?; - session.cancel_token = Some(cancel_token.clone()); - session.agent.clone() - }; + let agent = self + .get_session_agent(&session_id, Some(cancel_token.clone())) + .await?; let user_message = self.convert_acp_prompt_to_message(args.prompt); @@ -1107,13 +1146,7 @@ impl GooseAcpAgent { sacp::Error::internal_error().data(format!("Failed to create provider: {}", e)) })?; - let agent = { - let sessions = self.sessions.lock().await; - let session = sessions.get(session_id).ok_or_else(|| { - sacp::Error::invalid_params().data(format!("Session not found: {}", session_id)) - })?; - session.agent.clone() - }; + let agent = self.get_session_agent(session_id, None).await?; agent .update_provider(provider, session_id) .await @@ -1124,6 +1157,28 @@ impl GooseAcpAgent { info!(session_id = %session_id, model_id = %model_id, "Model switched"); Ok(SetSessionModelResponse::new()) } + + async fn on_set_mode( + &self, + session_id: &str, + mode_id: &str, + ) -> Result { + let mode = mode_id.parse::().map_err(|_| { + sacp::Error::invalid_params().data(format!("Invalid mode: {}", mode_id)) + })?; + + self.get_session_agent(session_id, None).await?; + + // Reject mode changes until per-session persistence lands (#7603) + if mode != self.goose_mode { + return Err(sacp::Error::invalid_params().data(format!( + "Mode change not supported: session is {}, requested {}", + self.goose_mode, mode_id + ))); + } + + Ok(SetSessionModeResponse::new()) + } } #[custom_methods] @@ -1135,7 +1190,7 @@ impl GooseAcpAgent { ) -> Result { let config: ExtensionConfig = serde_json::from_value(req.config) .map_err(|e| sacp::Error::invalid_params().data(format!("bad config: {e}")))?; - let agent = self.get_agent_for_session(&req.session_id).await?; + let agent = self.get_session_agent(&req.session_id, None).await?; agent .add_extension(config, &req.session_id) .await @@ -1148,7 +1203,7 @@ impl GooseAcpAgent { &self, req: RemoveExtensionRequest, ) -> Result { - let agent = self.get_agent_for_session(&req.session_id).await?; + let agent = self.get_session_agent(&req.session_id, None).await?; agent .remove_extension(&req.name, &req.session_id) .await @@ -1158,7 +1213,7 @@ impl GooseAcpAgent { #[custom_method("tools")] async fn on_get_tools(&self, req: GetToolsRequest) -> Result { - let agent = self.get_agent_for_session(&req.session_id).await?; + let agent = self.get_session_agent(&req.session_id, None).await?; let tools = agent.list_tools(&req.session_id, None).await; let tools_json = tools .into_iter() @@ -1173,7 +1228,7 @@ impl GooseAcpAgent { &self, req: ReadResourceRequest, ) -> Result { - let agent = self.get_agent_for_session(&req.session_id).await?; + let agent = self.get_session_agent(&req.session_id, None).await?; let cancel_token = CancellationToken::new(); let result = agent .extension_manager @@ -1310,17 +1365,6 @@ impl GooseAcpAgent { warnings, }) } - - async fn get_agent_for_session(&self, session_id: &str) -> Result, sacp::Error> { - self.sessions - .lock() - .await - .get(session_id) - .map(|s| Arc::clone(&s.agent)) - .ok_or_else(|| { - sacp::Error::invalid_params().data(format!("no active session: {session_id}")) - }) - } } pub struct GooseAcpHandler { @@ -1394,12 +1438,43 @@ impl JrMessageHandler for GooseAcpHandler { .if_notification(|notif: CancelNotification| async { agent.on_cancel(notif).await }) .await // Handle methods not yet in the sacp typed API. - // - session/set_model: typed support pending in sacp + // - session/set_model, session/set_mode: typed support pending in sacp // - _: custom requests that will eventually route to goose-server .otherwise({ let agent = agent.clone(); + let cx = cx.clone(); |message: MessageCx| async move { match message { + MessageCx::Request(req, request_cx) + if req.method == "session/set_mode" => + { + let params: SetSessionModeRequest = + serde_json::from_value(req.params).map_err(|e| { + sacp::Error::invalid_params().data(e.to_string()) + })?; + let session_id = params.session_id.clone(); + let mode_id = params.mode_id.clone(); + match agent.on_set_mode(&session_id.0, &mode_id.0).await { + Ok(resp) => { + let json = serde_json::to_value(resp).map_err(|e| { + sacp::Error::internal_error().data(e.to_string()) + })?; + // Notify before responding so clients see the mode + // update before block_task unblocks (serial dispatch). + cx.send_notification(SessionNotification::new( + session_id, + SessionUpdate::CurrentModeUpdate( + CurrentModeUpdate::new(mode_id), + ), + ))?; + request_cx.respond(json)?; + } + Err(e) => { + request_cx.respond_with_error(e)?; + } + } + Ok(()) + } MessageCx::Request(req, request_cx) if req.method == "session/set_model" => { @@ -1488,7 +1563,8 @@ mod tests { use rmcp::model::{CallToolRequestParams, Content as RmcpContent}; use sacp::schema::{ EnvVariable, HttpHeader, McpServer, McpServerHttp, McpServerSse, McpServerStdio, - PermissionOptionId, ResourceLink, SelectedPermissionOutcome, + PermissionOptionId, ResourceLink, SelectedPermissionOutcome, SessionMode, SessionModeId, + SessionModeState, }; use std::io::Write; use std::path::PathBuf; @@ -1841,4 +1917,25 @@ print(\"hello, world\") extract_locations_from_meta(&response) .map(|locs| locs.into_iter().map(|loc| (loc.path, loc.line)).collect()) } + + #[test] + fn test_build_mode_state() { + let state = build_mode_state(GooseMode::Auto).unwrap(); + assert_eq!( + state, + SessionModeState::new( + SessionModeId::new("auto"), + vec![ + SessionMode::new(SessionModeId::new("auto"), "auto") + .description("Automatically approve tool calls"), + SessionMode::new(SessionModeId::new("approve"), "approve") + .description("Ask before every tool call"), + SessionMode::new(SessionModeId::new("smart_approve"), "smart_approve") + .description("Ask only for sensitive tool calls"), + SessionMode::new(SessionModeId::new("chat"), "chat") + .description("Chat only, no tool calls"), + ], + ) + ); + } } diff --git a/crates/goose-acp/tests/common_tests/mod.rs b/crates/goose-acp/tests/common_tests/mod.rs index a9f537af..2e52e1a2 100644 --- a/crates/goose-acp/tests/common_tests/mod.rs +++ b/crates/goose-acp/tests/common_tests/mod.rs @@ -5,21 +5,22 @@ #[path = "../fixtures/mod.rs"] pub mod fixtures; use fixtures::{ - Connection, FsFixture, OpenAiFixture, PermissionDecision, Session, TestConnectionConfig, + Connection, FsFixture, OpenAiFixture, PermissionDecision, Session, SessionResult, + TestConnectionConfig, }; use fs_err as fs; use goose::config::base::CONFIG_YAML_NAME; use goose::config::GooseMode; use goose::providers::provider_registry::ProviderConstructor; -use goose_test_support::{ExpectedSessionId, McpFixture, FAKE_CODE, TEST_IMAGE_B64, TEST_MODEL}; -use sacp::schema::{McpServer, McpServerHttp, ModelId, ToolCallStatus}; +use goose_test_support::{McpFixture, FAKE_CODE, TEST_IMAGE_B64, TEST_MODEL}; +use sacp::schema::{McpServer, McpServerHttp, ModelId, SessionModeId, ToolCallStatus}; use std::sync::Arc; pub async fn run_config_mcp() { let temp_dir = tempfile::tempdir().unwrap(); - let expected_session_id = ExpectedSessionId::default(); + let expected_session_id = C::expected_session_id(); let prompt = "Use the get_code tool and output only its result."; - let mcp = McpFixture::new(Some(expected_session_id.clone())).await; + let mcp = McpFixture::new(expected_session_id.clone()).await; let config_yaml = format!( "GOOSE_MODEL: {TEST_MODEL}\nGOOSE_PROVIDER: openai\nextensions:\n mcp-fixture:\n enabled: true\n type: streamable_http\n name: mcp-fixture\n description: MCP fixture\n uri: \"{}\"\n", @@ -48,8 +49,8 @@ pub async fn run_config_mcp() { }; let mut conn = C::new(config, openai).await; - let (mut session, _) = conn.new_session().await; - expected_session_id.set(session.session_id().0.to_string()); + let SessionResult { mut session, .. } = conn.new_session().await; + expected_session_id.set(&session.session_id().0); let output = session.prompt(prompt, PermissionDecision::Cancel).await; assert_eq!(output.text, FAKE_CODE); @@ -64,7 +65,7 @@ pub async fn run_fs_read_text_file_true() { ); fs::write(temp_dir.path().join(CONFIG_YAML_NAME), config_yaml).unwrap(); - let expected_session_id = ExpectedSessionId::default(); + let expected_session_id = C::expected_session_id(); let prompt = "Use the read tool to read /tmp/test_acp_read.txt and output only its contents."; let openai = OpenAiFixture::new( vec![ @@ -88,8 +89,8 @@ pub async fn run_fs_read_text_file_true() { ..Default::default() }; let mut conn = C::new(config, openai).await; - let (mut session, _) = conn.new_session().await; - expected_session_id.set(session.session_id().0.to_string()); + let SessionResult { mut session, .. } = conn.new_session().await; + expected_session_id.set(&session.session_id().0); let output = session.prompt(prompt, PermissionDecision::Cancel).await; assert_eq!(output.text, "test-read-content-12345"); @@ -100,7 +101,7 @@ pub async fn run_fs_read_text_file_true() { pub async fn run_fs_write_text_file_false() { let _ = fs::remove_file("/tmp/test_acp_write.txt"); - let expected_session_id = ExpectedSessionId::default(); + let expected_session_id = C::expected_session_id(); let prompt = "Use the write tool to write 'test-write-content-67890' to /tmp/test_acp_write.txt"; let openai = OpenAiFixture::new( @@ -123,8 +124,8 @@ pub async fn run_fs_write_text_file_false() { ..Default::default() }; let mut conn = C::new(config, openai).await; - let (mut session, _) = conn.new_session().await; - expected_session_id.set(session.session_id().0.to_string()); + let SessionResult { mut session, .. } = conn.new_session().await; + expected_session_id.set(&session.session_id().0); let output = session.prompt(prompt, PermissionDecision::AllowOnce).await; assert!(!output.text.is_empty()); @@ -136,7 +137,7 @@ pub async fn run_fs_write_text_file_false() { } pub async fn run_fs_write_text_file_true() { - let expected_session_id = ExpectedSessionId::default(); + let expected_session_id = C::expected_session_id(); let prompt = "Use the write tool to write 'test-write-content-67890' to /tmp/test_acp_write.txt"; let openai = OpenAiFixture::new( @@ -163,8 +164,8 @@ pub async fn run_fs_write_text_file_true() { ..Default::default() }; let mut conn = C::new(config, openai).await; - let (mut session, _) = conn.new_session().await; - expected_session_id.set(session.session_id().0.to_string()); + let SessionResult { mut session, .. } = conn.new_session().await; + expected_session_id.set(&session.session_id().0); let output = session.prompt(prompt, PermissionDecision::AllowOnce).await; assert!(!output.text.is_empty()); @@ -176,7 +177,7 @@ pub async fn run_initialize_doesnt_hit_provider() { let provider_factory: ProviderConstructor = Arc::new(|_, _| Box::pin(async { Err(anyhow::anyhow!("no provider configured")) })); - let openai = OpenAiFixture::new(vec![], ExpectedSessionId::default()).await; + let openai = OpenAiFixture::new(vec![], C::expected_session_id()).await; let config = TestConnectionConfig { provider_factory: Some(provider_factory), ..Default::default() @@ -190,8 +191,67 @@ pub async fn run_initialize_doesnt_hit_provider() { .any(|m| &*m.id.0 == "goose-provider")); } +pub async fn run_load_mode() { + let temp_dir = tempfile::tempdir().unwrap(); + let expected_session_id = C::expected_session_id(); + let prompt = "Use the get_code tool and output only its result."; + let mcp = McpFixture::new(expected_session_id.clone()).await; + + let config_yaml = format!( + "GOOSE_MODEL: {TEST_MODEL}\nGOOSE_PROVIDER: openai\nextensions:\n mcp-fixture:\n enabled: true\n type: streamable_http\n name: mcp-fixture\n description: MCP fixture\n uri: \"{}\"\n", + mcp.url + ); + fs::write(temp_dir.path().join(CONFIG_YAML_NAME), config_yaml).unwrap(); + + let openai = OpenAiFixture::new( + vec![ + ( + prompt.to_string(), + include_str!("../test_data/openai_tool_call.txt"), + ), + ( + format!(r#""content":"{FAKE_CODE}""#), + include_str!("../test_data/openai_tool_result.txt"), + ), + ], + expected_session_id.clone(), + ) + .await; + + let config = TestConnectionConfig { + data_root: temp_dir.path().to_path_buf(), + ..Default::default() + }; + let mut conn = C::new(config, openai).await; + + let SessionResult { session, modes, .. } = conn.new_session().await; + assert_eq!( + modes.unwrap().current_mode_id, + SessionModeId::new(<&str>::from(GooseMode::default())) + ); + let session_id = session.session_id().0.to_string(); + conn.set_mode(&session_id, <&str>::from(GooseMode::Approve)) + .await + .unwrap(); + + let SessionResult { + session: mut loaded, + modes, + .. + } = conn.load_session(&session_id, vec![]).await; + assert_eq!( + modes.unwrap().current_mode_id, + SessionModeId::new(<&str>::from(GooseMode::Approve)) + ); + + // Approve mode + Cancel = permission denied → tool fails + expected_session_id.set(&loaded.session_id().0); + let output = loaded.prompt(prompt, PermissionDecision::Cancel).await; + assert_eq!(output.tool_status.unwrap(), ToolCallStatus::Failed); +} + pub async fn run_load_model() { - let expected_session_id = ExpectedSessionId::default(); + let expected_session_id = C::expected_session_id(); let openai = OpenAiFixture::new( vec![( r#""model":"o4-mini""#.into(), @@ -202,28 +262,88 @@ pub async fn run_load_model() { .await; let mut conn = C::new(TestConnectionConfig::default(), openai).await; - let (mut session, _) = conn.new_session().await; - expected_session_id.set(session.session_id().0.to_string()); + let SessionResult { mut session, .. } = conn.new_session().await; + expected_session_id.set(&session.session_id().0); - session.set_model("o4-mini").await; + let session_id = session.session_id().0.to_string(); + conn.set_model(&session_id, "o4-mini").await.unwrap(); let output = session .prompt("what is 1+1", PermissionDecision::Cancel) .await; assert_eq!(output.text, "2"); - let session_id = session.session_id().0.to_string(); - let (_, models) = conn.load_session(&session_id, vec![]).await; + let SessionResult { models, .. } = conn.load_session(&session_id, vec![]).await; assert_eq!(&*models.unwrap().current_model_id.0, "o4-mini"); } +pub async fn run_mode_set() { + let temp_dir = tempfile::tempdir().unwrap(); + let expected_session_id = C::expected_session_id(); + let prompt = "Use the get_code tool and output only its result."; + let mcp = McpFixture::new(expected_session_id.clone()).await; + + let config_yaml = format!( + "GOOSE_MODEL: {TEST_MODEL}\nGOOSE_PROVIDER: openai\nextensions:\n mcp-fixture:\n enabled: true\n type: streamable_http\n name: mcp-fixture\n description: MCP fixture\n uri: \"{}\"\n", + mcp.url + ); + fs::write(temp_dir.path().join(CONFIG_YAML_NAME), config_yaml).unwrap(); + + let openai = OpenAiFixture::new( + vec![ + ( + prompt.to_string(), + include_str!("../test_data/openai_tool_call.txt"), + ), + ( + format!(r#""content":"{FAKE_CODE}""#), + include_str!("../test_data/openai_tool_result.txt"), + ), + ], + expected_session_id.clone(), + ) + .await; + + let config = TestConnectionConfig { + data_root: temp_dir.path().to_path_buf(), + ..Default::default() + }; + let mut conn = C::new(config, openai).await; + + let SessionResult { + session: mut session_a, + .. + } = conn.new_session().await; + + let SessionResult { + session: mut session_b, + .. + } = conn.new_session().await; + conn.set_mode(&session_b.session_id().0, <&str>::from(GooseMode::Approve)) + .await + .unwrap(); + + // Approve mode + Cancel = permission denied → tool fails + expected_session_id.set(&session_b.session_id().0); + let output = session_b.prompt(prompt, PermissionDecision::Cancel).await; + assert_eq!(output.tool_status.unwrap(), ToolCallStatus::Failed); + + // Auto mode ignores Cancel — tool succeeds without permission prompt + conn.reset_openai(); + expected_session_id.set(&session_a.session_id().0); + let output = session_a.prompt(prompt, PermissionDecision::Cancel).await; + assert_eq!(output.text, FAKE_CODE); +} + pub async fn run_model_list() { - let expected_session_id = ExpectedSessionId::default(); + let expected_session_id = C::expected_session_id(); let openai = OpenAiFixture::new(vec![], expected_session_id.clone()).await; let mut conn = C::new(TestConnectionConfig::default(), openai).await; - let (session, models) = conn.new_session().await; - expected_session_id.set(session.session_id().0.to_string()); + let SessionResult { + session, models, .. + } = conn.new_session().await; + expected_session_id.set(&session.session_id().0); let models = models.unwrap(); assert!(!models.available_models.is_empty()); @@ -231,7 +351,7 @@ pub async fn run_model_list() { } pub async fn run_model_set() { - let expected_session_id = ExpectedSessionId::default(); + let expected_session_id = C::expected_session_id(); let openai = OpenAiFixture::new( vec![ // Session B prompt with switched model @@ -252,21 +372,29 @@ pub async fn run_model_set() { let mut conn = C::new(TestConnectionConfig::default(), openai).await; // Session A: default model - let (mut session_a, _) = conn.new_session().await; + let SessionResult { + session: mut session_a, + .. + } = conn.new_session().await; // Session B: switch to o4-mini - let (mut session_b, _) = conn.new_session().await; - session_b.set_model("o4-mini").await; + let SessionResult { + session: mut session_b, + .. + } = conn.new_session().await; + conn.set_model(&session_b.session_id().0, "o4-mini") + .await + .unwrap(); // Prompt B — expects o4-mini - expected_session_id.set(session_b.session_id().0.to_string()); + expected_session_id.set(&session_b.session_id().0); let output = session_b .prompt("what is 1+1", PermissionDecision::Cancel) .await; assert_eq!(output.text, "2"); // Prompt A — expects default TEST_MODEL (proves sessions are independent) - expected_session_id.set(session_a.session_id().0.to_string()); + expected_session_id.set(&session_a.session_id().0); let output = session_a .prompt("what is 1+1", PermissionDecision::Cancel) .await; @@ -292,8 +420,8 @@ pub async fn run_permission_persistence() { let temp_dir = tempfile::tempdir().unwrap(); let prompt = "Use the get_code tool and output only its result."; - let expected_session_id = ExpectedSessionId::default(); - let mcp = McpFixture::new(Some(expected_session_id.clone())).await; + let expected_session_id = C::expected_session_id(); + let mcp = McpFixture::new(expected_session_id.clone()).await; let openai = OpenAiFixture::new( vec![ ( @@ -317,8 +445,8 @@ pub async fn run_permission_persistence() { }; let mut conn = C::new(config, openai).await; - let (mut session, _) = conn.new_session().await; - expected_session_id.set(session.session_id().0.to_string()); + let SessionResult { mut session, .. } = conn.new_session().await; + expected_session_id.set(&session.session_id().0); for (decision, expected_status, expected_yaml) in cases { conn.reset_openai(); @@ -336,7 +464,7 @@ pub async fn run_permission_persistence() { } pub async fn run_prompt_basic() { - let expected_session_id = ExpectedSessionId::default(); + let expected_session_id = C::expected_session_id(); let openai = OpenAiFixture::new( vec![( r#"\nwhat is 1+1""#.into(), @@ -347,8 +475,8 @@ pub async fn run_prompt_basic() { .await; let mut conn = C::new(TestConnectionConfig::default(), openai).await; - let (mut session, _) = conn.new_session().await; - expected_session_id.set(session.session_id().0.to_string()); + let SessionResult { mut session, .. } = conn.new_session().await; + expected_session_id.set(&session.session_id().0); let output = session .prompt("what is 1+1", PermissionDecision::Cancel) @@ -358,10 +486,10 @@ pub async fn run_prompt_basic() { } pub async fn run_prompt_codemode() { - let expected_session_id = ExpectedSessionId::default(); + let expected_session_id = C::expected_session_id(); let prompt = "Search for getCode and write tools. Use them to save the code to /tmp/result.txt."; - let mcp = McpFixture::new(Some(expected_session_id.clone())).await; + let mcp = McpFixture::new(expected_session_id.clone()).await; let openai = OpenAiFixture::new( vec![ ( @@ -390,8 +518,8 @@ pub async fn run_prompt_codemode() { let _ = fs::remove_file("/tmp/result.txt"); let mut conn = C::new(config, openai).await; - let (mut session, _) = conn.new_session().await; - expected_session_id.set(session.session_id().0.to_string()); + let SessionResult { mut session, .. } = conn.new_session().await; + expected_session_id.set(&session.session_id().0); let output = session.prompt(prompt, PermissionDecision::Cancel).await; if matches!(output.tool_status, Some(ToolCallStatus::Failed)) || output.text.contains("error") { @@ -404,8 +532,8 @@ pub async fn run_prompt_codemode() { } pub async fn run_prompt_image() { - let expected_session_id = ExpectedSessionId::default(); - let mcp = McpFixture::new(Some(expected_session_id.clone())).await; + let expected_session_id = C::expected_session_id(); + let mcp = McpFixture::new(expected_session_id.clone()).await; let openai = OpenAiFixture::new( vec![ ( @@ -427,8 +555,8 @@ pub async fn run_prompt_image() { ..Default::default() }; let mut conn = C::new(config, openai).await; - let (mut session, _) = conn.new_session().await; - expected_session_id.set(session.session_id().0.to_string()); + let SessionResult { mut session, .. } = conn.new_session().await; + expected_session_id.set(&session.session_id().0); let output = session .prompt( @@ -441,7 +569,7 @@ pub async fn run_prompt_image() { } pub async fn run_prompt_image_attachment() { - let expected_session_id = ExpectedSessionId::default(); + let expected_session_id = C::expected_session_id(); let openai = OpenAiFixture::new( vec![( r#""type":"image_url""#.into(), @@ -452,8 +580,8 @@ pub async fn run_prompt_image_attachment() { .await; let mut conn = C::new(TestConnectionConfig::default(), openai).await; - let (mut session, _) = conn.new_session().await; - expected_session_id.set(session.session_id().0.to_string()); + let SessionResult { mut session, .. } = conn.new_session().await; + expected_session_id.set(&session.session_id().0); let output = session .prompt_with_image( @@ -468,9 +596,9 @@ pub async fn run_prompt_image_attachment() { } pub async fn run_load_session_mcp() { - let expected_session_id = ExpectedSessionId::default(); + let expected_session_id = C::expected_session_id(); let prompt = "Use the get_code tool and output only its result."; - let mcp = McpFixture::new(Some(expected_session_id.clone())).await; + let mcp = McpFixture::new(expected_session_id.clone()).await; let mcp_url = mcp.url.clone(); // Two rounds of tool call + tool result: one for new session, one for loaded session. @@ -504,8 +632,8 @@ pub async fn run_load_session_mcp() { ..Default::default() }; let mut conn = C::new(config, openai).await; - let (mut session, _) = conn.new_session().await; - expected_session_id.set(session.session_id().0.to_string()); + let SessionResult { mut session, .. } = conn.new_session().await; + expected_session_id.set(&session.session_id().0); // First prompt: tool should work in the new session. let output = session.prompt(prompt, PermissionDecision::Cancel).await; @@ -513,7 +641,10 @@ pub async fn run_load_session_mcp() { // Load the same session with MCP servers re-specified. let session_id = session.session_id().0.to_string(); - let (mut loaded_session, _) = conn.load_session(&session_id, mcp_servers).await; + let SessionResult { + session: mut loaded_session, + .. + } = conn.load_session(&session_id, mcp_servers).await; // Second prompt: tool should work in the loaded session. let output = loaded_session @@ -523,8 +654,8 @@ pub async fn run_load_session_mcp() { } pub async fn run_prompt_mcp() { - let expected_session_id = ExpectedSessionId::default(); - let mcp = McpFixture::new(Some(expected_session_id.clone())).await; + let expected_session_id = C::expected_session_id(); + let mcp = McpFixture::new(expected_session_id.clone()).await; let openai = OpenAiFixture::new( vec![ ( @@ -545,8 +676,8 @@ pub async fn run_prompt_mcp() { ..Default::default() }; let mut conn = C::new(config, openai).await; - let (mut session, _) = conn.new_session().await; - expected_session_id.set(session.session_id().0.to_string()); + let SessionResult { mut session, .. } = conn.new_session().await; + expected_session_id.set(&session.session_id().0); let output = session .prompt( @@ -557,3 +688,46 @@ pub async fn run_prompt_mcp() { assert_eq!(output.text, FAKE_CODE); expected_session_id.assert_matches(&session.session_id().0); } + +pub async fn run_mode_set_error( + mode_id: &str, + session_id_override: Option<&str>, + expected: sacp::Error, +) { + let openai = OpenAiFixture::new(vec![], C::expected_session_id()).await; + let mut conn = C::new(TestConnectionConfig::default(), openai).await; + let SessionResult { session, .. } = conn.new_session().await; + + let target_session_id = session_id_override + .map(str::to_string) + .unwrap_or_else(|| session.session_id().0.to_string()); + + let err = conn + .set_mode(&target_session_id, mode_id) + .await + .unwrap_err(); + + let sacp_err = err.downcast::().unwrap(); + assert_eq!(sacp_err, expected); +} + +#[macro_export] +macro_rules! tests_mode_set_error { + ($conn:ty) => { + #[test_case::test_case("not_a_mode", None, sacp::Error::invalid_params().data("Invalid mode: not_a_mode") ; "invalid mode")] + #[test_case::test_case("auto", Some("nonexistent-session-id"), sacp::Error::invalid_params().data("Session not found: nonexistent-session-id") ; "session not found")] + #[test_case::test_case("approve", None, sacp::Error::invalid_params().data("Mode change not supported: session is auto, requested approve") ; "mode change rejected")] + fn test_mode_set_error( + mode_id: &'static str, + session_id: Option<&'static str>, + expected: sacp::Error, + ) { + common_tests::fixtures::run_test(async move { + common_tests::run_mode_set_error::<$conn>( + mode_id, session_id, expected, + ) + .await + }); + } + }; +} diff --git a/crates/goose-acp/tests/custom_requests_test.rs b/crates/goose-acp/tests/custom_requests_test.rs index 0f7c7678..01606317 100644 --- a/crates/goose-acp/tests/custom_requests_test.rs +++ b/crates/goose-acp/tests/custom_requests_test.rs @@ -2,8 +2,9 @@ mod common_tests; use common_tests::fixtures::server::ClientToAgentConnection; -use common_tests::fixtures::{run_test, Connection, Session, TestConnectionConfig}; -use goose_test_support::ExpectedSessionId; +use common_tests::fixtures::{run_test, Connection, Session, SessionResult, TestConnectionConfig}; +use goose_test_support::EnforceSessionId; +use std::sync::Arc; use common_tests::fixtures::OpenAiFixture; @@ -20,10 +21,10 @@ async fn send_custom( #[test] fn test_custom_session_list() { run_test(async { - let openai = OpenAiFixture::new(vec![], ExpectedSessionId::default()).await; + let openai = OpenAiFixture::new(vec![], Arc::new(EnforceSessionId::default())).await; let mut conn = ClientToAgentConnection::new(TestConnectionConfig::default(), openai).await; - let (session, _models) = conn.new_session().await; + let SessionResult { session, .. } = conn.new_session().await; let session_id = session.session_id().0.clone(); // Verify the session exists via _session/get @@ -61,10 +62,10 @@ fn test_custom_session_list() { #[test] fn test_custom_session_get() { run_test(async { - let openai = OpenAiFixture::new(vec![], ExpectedSessionId::default()).await; + let openai = OpenAiFixture::new(vec![], Arc::new(EnforceSessionId::default())).await; let mut conn = ClientToAgentConnection::new(TestConnectionConfig::default(), openai).await; - let (session, _models) = conn.new_session().await; + let SessionResult { session, .. } = conn.new_session().await; let session_id = session.session_id().0.clone(); let result = send_custom( @@ -89,10 +90,10 @@ fn test_custom_session_get() { #[test] fn test_custom_session_delete() { run_test(async { - let openai = OpenAiFixture::new(vec![], ExpectedSessionId::default()).await; + let openai = OpenAiFixture::new(vec![], Arc::new(EnforceSessionId::default())).await; let mut conn = ClientToAgentConnection::new(TestConnectionConfig::default(), openai).await; - let (session, _models) = conn.new_session().await; + let SessionResult { session, .. } = conn.new_session().await; let session_id = session.session_id().0.clone(); let result = send_custom( @@ -116,10 +117,10 @@ fn test_custom_session_delete() { #[test] fn test_custom_get_tools() { run_test(async { - let openai = OpenAiFixture::new(vec![], ExpectedSessionId::default()).await; + let openai = OpenAiFixture::new(vec![], Arc::new(EnforceSessionId::default())).await; let mut conn = ClientToAgentConnection::new(TestConnectionConfig::default(), openai).await; - let (session, _models) = conn.new_session().await; + let SessionResult { session, .. } = conn.new_session().await; let session_id = session.session_id().0.clone(); let result = send_custom( @@ -139,7 +140,7 @@ fn test_custom_get_tools() { #[test] fn test_custom_get_extensions() { run_test(async { - let openai = OpenAiFixture::new(vec![], ExpectedSessionId::default()).await; + let openai = OpenAiFixture::new(vec![], Arc::new(EnforceSessionId::default())).await; let conn = ClientToAgentConnection::new(TestConnectionConfig::default(), openai).await; let result = @@ -161,7 +162,7 @@ fn test_custom_get_extensions() { #[test] fn test_custom_unknown_method() { run_test(async { - let openai = OpenAiFixture::new(vec![], ExpectedSessionId::default()).await; + let openai = OpenAiFixture::new(vec![], Arc::new(EnforceSessionId::default())).await; let conn = ClientToAgentConnection::new(TestConnectionConfig::default(), openai).await; let result = send_custom(conn.cx(), "_unknown/method", serde_json::json!({})).await; diff --git a/crates/goose-acp/tests/fixtures/mod.rs b/crates/goose-acp/tests/fixtures/mod.rs index 47105032..fef698b6 100644 --- a/crates/goose-acp/tests/fixtures/mod.rs +++ b/crates/goose-acp/tests/fixtures/mod.rs @@ -14,8 +14,8 @@ use goose::session_context::SESSION_ID_HEADER; use goose_acp::server::{serve, GooseAcpAgent}; use goose_test_support::{ExpectedSessionId, TEST_MODEL}; use sacp::schema::{ - AuthMethod, McpServer, ReadTextFileRequest, ReadTextFileResponse, SessionModelState, - ToolCallStatus, WriteTextFileRequest, WriteTextFileResponse, + AuthMethod, McpServer, ReadTextFileRequest, ReadTextFileResponse, SessionModeState, + SessionModelState, ToolCallStatus, WriteTextFileRequest, WriteTextFileResponse, }; use std::collections::VecDeque; use std::future::Future; @@ -38,7 +38,7 @@ impl OpenAiFixture { /// On mismatch, returns 417 of the diff in OpenAI error format. pub async fn new( exchanges: Vec<(String, &'static str)>, - expected_session_id: ExpectedSessionId, + expected_session_id: Arc, ) -> Self { let mock_server = MockServer::start().await; let queue = Arc::new(Mutex::new(VecDeque::from(exchanges.clone()))); @@ -266,6 +266,12 @@ impl FsFixture { } } +pub struct SessionResult { + pub session: S, + pub models: Option, + pub modes: Option, +} + pub struct TestConnectionConfig { pub mcp_servers: Vec, pub builtins: Vec, @@ -281,7 +287,7 @@ impl Default for TestConnectionConfig { Self { mcp_servers: Vec::new(), builtins: Vec::new(), - goose_mode: GooseMode::Auto, + goose_mode: GooseMode::default(), data_root: PathBuf::new(), provider_factory: None, read_text_file: None, @@ -294,13 +300,16 @@ impl Default for TestConnectionConfig { pub trait Connection: Sized { type Session: Session; + fn expected_session_id() -> Arc; async fn new(config: TestConnectionConfig, openai: OpenAiFixture) -> Self; - async fn new_session(&mut self) -> (Self::Session, Option); + async fn new_session(&mut self) -> SessionResult; async fn load_session( &mut self, session_id: &str, mcp_servers: Vec, - ) -> (Self::Session, Option); + ) -> SessionResult; + async fn set_mode(&self, session_id: &str, mode_id: &str) -> anyhow::Result<()>; + async fn set_model(&self, session_id: &str, model_id: &str) -> anyhow::Result<()>; fn auth_methods(&self) -> &[AuthMethod]; fn reset_openai(&self); fn reset_permissions(&self); @@ -317,7 +326,6 @@ pub trait Session { mime_type: &str, decision: PermissionDecision, ) -> TestOutput; - async fn set_model(&self, model_id: &str); } #[allow(dead_code)] diff --git a/crates/goose-acp/tests/fixtures/provider.rs b/crates/goose-acp/tests/fixtures/provider.rs index 7256b8cd..608b2d91 100644 --- a/crates/goose-acp/tests/fixtures/provider.rs +++ b/crates/goose-acp/tests/fixtures/provider.rs @@ -1,18 +1,21 @@ use super::{ spawn_acp_server_in_process, Connection, OpenAiFixture, PermissionDecision, Session, - TestConnectionConfig, TestOutput, + SessionResult, TestConnectionConfig, TestOutput, }; use async_trait::async_trait; use futures::StreamExt; use goose::acp::{AcpProvider, AcpProviderConfig, PermissionMapping}; +use goose::config::goose_mode::GooseMode; use goose::config::PermissionManager; use goose::conversation::message::{ActionRequiredData, Message, MessageContent}; use goose::model::ModelConfig; use goose::permission::permission_confirmation::PrincipalType; use goose::permission::{Permission, PermissionConfirmation}; use goose::providers::base::Provider; -use goose_test_support::TEST_MODEL; -use sacp::schema::{AuthMethod, McpServer, SessionModelState, ToolCallStatus}; +use goose::providers::errors::ProviderError; +use goose_test_support::{ExpectedSessionId, IgnoreSessionId, TEST_MODEL}; +use sacp::schema::{AuthMethod, McpServer, ToolCallStatus}; +use std::str::FromStr; use std::sync::Arc; use tokio::sync::Mutex; @@ -29,14 +32,13 @@ pub struct ClientToProviderConnection { #[allow(dead_code)] pub struct ClientToProviderSession { provider: Arc>, - acp_session_id: sacp::schema::SessionId, - session_id: String, + session_id: sacp::schema::SessionId, } impl ClientToProviderSession { #[allow(dead_code)] async fn send_message(&mut self, message: Message, decision: PermissionDecision) -> TestOutput { - let session_id = self.session_id.clone(); + let session_id = self.session_id.0.clone(); let provider = self.provider.lock().await; let model_config = provider.get_model_config(); let mut stream = provider @@ -101,6 +103,10 @@ impl ClientToProviderSession { impl Connection for ClientToProviderConnection { type Session = ClientToProviderSession; + fn expected_session_id() -> Arc { + Arc::new(IgnoreSessionId) + } + async fn new(config: TestConnectionConfig, openai: OpenAiFixture) -> Self { let (data_root, temp_dir) = match config.data_root.as_os_str().is_empty() { true => { @@ -156,7 +162,7 @@ impl Connection for ClientToProviderConnection { } } - async fn new_session(&mut self) -> (ClientToProviderSession, Option) { + async fn new_session(&mut self) -> SessionResult { // Tests like run_model_set call new_session() multiple times on the same // connection, so each needs a distinct key to avoid returning a cached session. self.session_counter += 1; @@ -171,20 +177,51 @@ impl Connection for ClientToProviderConnection { let session = ClientToProviderSession { provider: Arc::clone(&self.provider), - acp_session_id: response.session_id, - session_id: goose_id, + session_id: sacp::schema::SessionId::new(goose_id), }; - (session, response.models) + SessionResult { + session, + models: response.models, + modes: response.modes, + } } async fn load_session( &mut self, _session_id: &str, _mcp_servers: Vec, - ) -> (ClientToProviderSession, Option) { + ) -> SessionResult { unimplemented!("TODO: implement load_session in ACP provider") } + async fn set_mode(&self, session_id: &str, mode_id: &str) -> anyhow::Result<()> { + let mode = GooseMode::from_str(mode_id).map_err(|_| { + sacp::Error::invalid_params().data(format!("Invalid mode: {}", mode_id)) + })?; + self.provider + .lock() + .await + .update_mode(session_id, mode) + .await + .map_err(|e| match e { + ProviderError::RequestFailed(msg) => sacp::Error::invalid_params().data(msg), + other => sacp::Error::internal_error().data(other.to_string()), + })?; + Ok(()) + } + + async fn set_model(&self, session_id: &str, model_id: &str) -> anyhow::Result<()> { + let provider = self.provider.lock().await; + let response = provider.ensure_session(Some(session_id)).await?; + provider + .send_untyped( + "session/set_model", + serde_json::json!({ "sessionId": response.session_id, "modelId": model_id }), + ) + .await?; + Ok(()) + } + fn auth_methods(&self) -> &[AuthMethod] { &self.auth_methods } @@ -201,7 +238,7 @@ impl Connection for ClientToProviderConnection { #[async_trait] impl Session for ClientToProviderSession { fn session_id(&self) -> &sacp::schema::SessionId { - &self.acp_session_id + &self.session_id } async fn prompt(&mut self, prompt: &str, decision: PermissionDecision) -> TestOutput { @@ -221,13 +258,4 @@ impl Session for ClientToProviderSession { .with_text(prompt); self.send_message(message, decision).await } - - async fn set_model(&self, model_id: &str) { - self.provider - .lock() - .await - .set_model(&self.acp_session_id, model_id) - .await - .unwrap(); - } } diff --git a/crates/goose-acp/tests/fixtures/server.rs b/crates/goose-acp/tests/fixtures/server.rs index 95c6319d..0c4180c0 100644 --- a/crates/goose-acp/tests/fixtures/server.rs +++ b/crates/goose-acp/tests/fixtures/server.rs @@ -1,15 +1,15 @@ use super::{ map_permission_response, spawn_acp_server_in_process, Connection, PermissionDecision, - PermissionMapping, Session, TestConnectionConfig, TestOutput, + PermissionMapping, Session, SessionResult, TestConnectionConfig, TestOutput, }; use async_trait::async_trait; use goose::config::PermissionManager; +use goose_test_support::{EnforceSessionId, ExpectedSessionId}; use sacp::schema::{ AuthMethod, ClientCapabilities, ContentBlock, FileSystemCapability, ImageContent, InitializeRequest, LoadSessionRequest, McpServer, NewSessionRequest, PromptRequest, - ProtocolVersion, ReadTextFileRequest, RequestPermissionRequest, SessionModelState, - SessionNotification, SessionUpdate, StopReason, TextContent, ToolCallStatus, - WriteTextFileRequest, + ProtocolVersion, ReadTextFileRequest, RequestPermissionRequest, SessionNotification, + SessionUpdate, StopReason, TextContent, ToolCallStatus, WriteTextFileRequest, }; use sacp::{ClientToAgent, JrConnectionCx}; use std::sync::{Arc, Mutex}; @@ -84,6 +84,10 @@ impl ClientToAgentConnection { impl Connection for ClientToAgentConnection { type Session = ClientToAgentSession; + fn expected_session_id() -> Arc { + Arc::new(EnforceSessionId::default()) + } + async fn new(config: TestConnectionConfig, openai: super::OpenAiFixture) -> Self { let (data_root, temp_dir) = match config.data_root.as_os_str().is_empty() { true => { @@ -228,7 +232,7 @@ impl Connection for ClientToAgentConnection { } } - async fn new_session(&mut self) -> (ClientToAgentSession, Option) { + async fn new_session(&mut self) -> SessionResult { let work_dir = tempfile::tempdir().unwrap(); let mcp_servers = std::mem::take(&mut self.pending_mcp_servers); let response = self @@ -244,14 +248,18 @@ impl Connection for ClientToAgentConnection { permission: self.permission.clone(), notify: self.notify.clone(), }; - (session, response.models) + SessionResult { + session, + models: response.models, + modes: response.modes, + } } async fn load_session( &mut self, session_id: &str, mcp_servers: Vec, - ) -> (ClientToAgentSession, Option) { + ) -> SessionResult { self.updates.lock().unwrap().clear(); let work_dir = tempfile::tempdir().unwrap(); let session_id = sacp::schema::SessionId::new(session_id.to_string()); @@ -271,7 +279,37 @@ impl Connection for ClientToAgentConnection { permission: self.permission.clone(), notify: self.notify.clone(), }; - (session, response.models) + SessionResult { + session, + models: response.models, + modes: response.modes, + } + } + + async fn set_mode(&self, session_id: &str, mode_id: &str) -> anyhow::Result<()> { + let msg = sacp::UntypedMessage::new( + "session/set_mode", + serde_json::json!({ "sessionId": session_id, "modeId": mode_id }), + )?; + self.cx + .send_request(msg) + .block_task() + .await + .map(|_| ()) + .map_err(|e| e.into()) + } + + async fn set_model(&self, session_id: &str, model_id: &str) -> anyhow::Result<()> { + let msg = sacp::UntypedMessage::new( + "session/set_model", + serde_json::json!({ "sessionId": session_id, "modelId": model_id }), + )?; + self.cx + .send_request(msg) + .block_task() + .await + .map(|_| ()) + .map_err(|e| e.into()) } fn auth_methods(&self) -> &[AuthMethod] { @@ -314,19 +352,6 @@ impl Session for ClientToAgentSession { ) .await } - - // HACK: sacp doesn't support session/set_model yet, so we send it as untyped JSON. - async fn set_model(&self, model_id: &str) { - let msg = sacp::UntypedMessage::new( - "session/set_model", - serde_json::json!({ - "sessionId": self.session_id.0, - "modelId": model_id - }), - ) - .unwrap(); - self.cx.send_request(msg).block_task().await.unwrap(); - } } fn collect_agent_text(updates: &Arc>>) -> String { diff --git a/crates/goose-acp/tests/provider_test.rs b/crates/goose-acp/tests/provider_test.rs index 5b1da7be..34f0435a 100644 --- a/crates/goose-acp/tests/provider_test.rs +++ b/crates/goose-acp/tests/provider_test.rs @@ -5,12 +5,14 @@ use common_tests::fixtures::provider::ClientToProviderConnection; use common_tests::fixtures::run_test; use common_tests::{ run_config_mcp, run_fs_read_text_file_true, run_fs_write_text_file_false, - run_fs_write_text_file_true, run_initialize_doesnt_hit_provider, run_load_model, - run_load_session_mcp, run_model_list, run_model_set, run_permission_persistence, + run_fs_write_text_file_true, run_initialize_doesnt_hit_provider, run_load_mode, run_load_model, + run_load_session_mcp, run_mode_set, run_model_list, run_model_set, run_permission_persistence, run_prompt_basic, run_prompt_codemode, run_prompt_image, run_prompt_image_attachment, run_prompt_mcp, }; +tests_mode_set_error!(ClientToProviderConnection); + #[test] fn test_config_mcp() { run_test(async { run_config_mcp::().await }); @@ -38,6 +40,12 @@ fn test_initialize_doesnt_hit_provider() { run_test(async { run_initialize_doesnt_hit_provider::().await }); } +#[test] +#[ignore = "TODO: implement load_session in ACP provider"] +fn test_load_mode() { + run_test(async { run_load_mode::().await }); +} + #[test] #[ignore = "TODO: implement load_session in ACP provider"] fn test_load_model() { @@ -50,6 +58,12 @@ fn test_load_session_mcp() { run_test(async { run_load_session_mcp::().await }); } +#[test] +#[ignore = "TODO: on_set_mode is a no-op until mode is threaded per-session (#7603)"] +fn test_mode_set() { + run_test(async { run_mode_set::().await }); +} + #[test] fn test_model_list() { run_test(async { run_model_list::().await }); diff --git a/crates/goose-acp/tests/server_test.rs b/crates/goose-acp/tests/server_test.rs index 0f65f24f..c3a5ac72 100644 --- a/crates/goose-acp/tests/server_test.rs +++ b/crates/goose-acp/tests/server_test.rs @@ -3,12 +3,14 @@ use common_tests::fixtures::run_test; use common_tests::fixtures::server::ClientToAgentConnection; use common_tests::{ run_config_mcp, run_fs_read_text_file_true, run_fs_write_text_file_false, - run_fs_write_text_file_true, run_initialize_doesnt_hit_provider, run_load_model, - run_load_session_mcp, run_model_list, run_model_set, run_permission_persistence, + run_fs_write_text_file_true, run_initialize_doesnt_hit_provider, run_load_mode, run_load_model, + run_load_session_mcp, run_mode_set, run_model_list, run_model_set, run_permission_persistence, run_prompt_basic, run_prompt_codemode, run_prompt_image, run_prompt_image_attachment, run_prompt_mcp, }; +tests_mode_set_error!(ClientToAgentConnection); + #[test] fn test_config_mcp() { run_test(async { run_config_mcp::().await }); @@ -34,6 +36,12 @@ fn test_initialize_doesnt_hit_provider() { run_test(async { run_initialize_doesnt_hit_provider::().await }); } +#[test] +#[ignore = "TODO: on_set_mode is a no-op until mode is threaded per-session (#7603)"] +fn test_load_mode() { + run_test(async { run_load_mode::().await }); +} + #[test] fn test_load_model() { run_test(async { run_load_model::().await }); @@ -44,6 +52,12 @@ fn test_load_session_mcp() { run_test(async { run_load_session_mcp::().await }); } +#[test] +#[ignore = "TODO: on_set_mode is a no-op until mode is threaded per-session (#7603)"] +fn test_mode_set() { + run_test(async { run_mode_set::().await }); +} + #[test] fn test_model_list() { run_test(async { run_model_list::().await }); diff --git a/crates/goose-test-support/src/lib.rs b/crates/goose-test-support/src/lib.rs index 07cfcd7b..be80c4a8 100644 --- a/crates/goose-test-support/src/lib.rs +++ b/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, +}; diff --git a/crates/goose-test-support/src/mcp.rs b/crates/goose-test-support/src/mcp.rs index 519f57ac..230c9387 100644 --- a/crates/goose-test-support/src/mcp.rs +++ b/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 HasMeta for NotificationContext { struct ValidatingService { inner: S, - expected_session_id: ExpectedSessionId, + expected_session_id: Arc, } impl ValidatingService { - fn new(inner: S, expected_session_id: ExpectedSessionId) -> Self { + fn new(inner: S, expected_session_id: Arc) -> Self { Self { inner, expected_session_id, @@ -144,16 +146,13 @@ type McpServiceFactory = Box Result>, std::io::Error> + Send + Sync>; impl McpFixture { - pub async fn new(expected_session_id: Option) -> 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) -> 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, diff --git a/crates/goose-test-support/src/session.rs b/crates/goose-test-support/src/session.rs index 30e71290..c150e82c 100644 --- a/crates/goose-test-support/src/session.rs +++ b/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>, errors: Arc>>, } -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) { +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) {} +} diff --git a/crates/goose/src/acp/provider.rs b/crates/goose/src/acp/provider.rs index 461791af..44917a1f 100644 --- a/crates/goose/src/acp/provider.rs +++ b/crates/goose/src/acp/provider.rs @@ -13,6 +13,7 @@ use sacp::{ClientToAgent, JrConnectionCx}; use std::collections::{HashMap, HashSet}; use std::path::PathBuf; use std::process::Stdio; +use std::str::FromStr; use std::sync::{Arc, Mutex}; use tokio::process::{Child, Command}; use tokio::sync::{mpsc, oneshot, Mutex as TokioMutex}; @@ -43,10 +44,10 @@ enum ClientRequest { NewSession { response_tx: oneshot::Sender>, }, - SetModel { - session_id: SessionId, - model_id: String, - response_tx: oneshot::Sender>, + Untyped { + method: String, + params: serde_json::Value, + response_tx: oneshot::Sender>, }, Prompt { session_id: SessionId, @@ -77,7 +78,7 @@ enum AcpUpdate { pub struct AcpProvider { name: String, model: ModelConfig, - goose_mode: GooseMode, + goose_mode: Arc>, tx: mpsc::Sender, permission_mapping: PermissionMapping, rejected_tool_calls: Arc>>, @@ -107,8 +108,10 @@ impl AcpProvider { let (init_tx, init_rx) = oneshot::channel(); let permission_mapping = config.permission_mapping.clone(); let rejected_tool_calls = Arc::new(TokioMutex::new(HashSet::new())); + let goose_mode = Arc::new(Mutex::new(goose_mode)); - tokio::spawn(run_client_loop(config, rx, init_tx)); + let client_loop = AcpClientLoop::new(config, goose_mode.clone()); + tokio::spawn(client_loop.spawn(rx, init_tx)); let init_response = init_rx .await @@ -141,12 +144,11 @@ impl AcpProvider { let (init_tx, init_rx) = oneshot::channel(); let permission_mapping = config.permission_mapping.clone(); let rejected_tool_calls = Arc::new(TokioMutex::new(HashSet::new())); + let goose_mode = Arc::new(Mutex::new(goose_mode)); let transport = sacp::ByteStreams::new(write, read); - let init_tx = Arc::new(Mutex::new(Some(init_tx))); + let client_loop = AcpClientLoop::new(config, goose_mode.clone()); tokio::spawn(async move { - if let Err(e) = - run_protocol_loop_with_transport(config, transport, &mut rx, init_tx.clone()).await - { + if let Err(e) = client_loop.run(transport, &mut rx, init_tx).await { tracing::error!("ACP protocol error: {e}"); } }); @@ -169,7 +171,7 @@ impl AcpProvider { fn new_with_runtime( name: String, model: ModelConfig, - goose_mode: GooseMode, + goose_mode: Arc>, tx: mpsc::Sender, permission_mapping: PermissionMapping, rejected_tool_calls: Arc>>, @@ -201,19 +203,21 @@ impl AcpProvider { response_rx.await.context("ACP session/new cancelled")? } - pub async fn set_model(&self, session_id: &SessionId, model_id: &str) -> Result<()> { + pub async fn send_untyped( + &self, + method: &str, + params: serde_json::Value, + ) -> Result { let (response_tx, response_rx) = oneshot::channel(); self.tx - .send(ClientRequest::SetModel { - session_id: session_id.clone(), - model_id: model_id.to_string(), + .send(ClientRequest::Untyped { + method: method.to_string(), + params, response_tx, }) .await .context("ACP client is unavailable")?; - response_rx - .await - .context("ACP session/set_model cancelled")? + response_rx.await.context("ACP request cancelled")? } pub async fn handle_permission_confirmation( @@ -281,6 +285,32 @@ impl Provider for AcpProvider { self.model.clone() } + async fn update_mode(&self, session_id: &str, mode: GooseMode) -> Result<(), ProviderError> { + let _acp_session_id = self + .goose_to_acp_id + .lock() + .await + .get(session_id) + .map(|r| r.session_id.clone()) + .ok_or_else(|| { + ProviderError::RequestFailed(format!("Session not found: {session_id}")) + })?; + + let current = self + .goose_mode + .lock() + .map_err(|_| ProviderError::RequestFailed("Failed to read mode".into()))?; + + if mode != *current { + // TODO: "session/set_mode" when session-scoped mode lands (#7603) + return Err(ProviderError::RequestFailed(format!( + "Mode change not supported: session is {}, requested {}", + current, mode + ))); + } + Ok(()) + } + fn permission_routing(&self) -> PermissionRouting { PermissionRouting::ActionRequired } @@ -311,7 +341,10 @@ impl Provider for AcpProvider { let pending_confirmations = self.pending_confirmations.clone(); let rejected_tool_calls = self.rejected_tool_calls.clone(); let permission_mapping = self.permission_mapping.clone(); - let goose_mode = self.goose_mode; + let goose_mode = *self + .goose_mode + .lock() + .map_err(|_| ProviderError::RequestFailed("goose_mode lock poisoned".into()))?; let reject_all_tools = goose_mode == GooseMode::Chat; @@ -422,31 +455,161 @@ impl Drop for AcpProvider { } } -async fn run_client_loop( +struct AcpClientLoop { config: AcpProviderConfig, - mut rx: mpsc::Receiver, - init_tx: oneshot::Sender>, -) { - let init_tx = Arc::new(Mutex::new(Some(init_tx))); + goose_mode: Arc>, + prompt_response_tx: Arc>>>, +} - let child = match spawn_acp_process(&config).await { - Ok(c) => c, - Err(e) => { - let message = e.to_string(); - send_init_result(&init_tx, Err(anyhow::anyhow!(message.clone()))); - tracing::error!("failed to spawn ACP process: {message}"); - return; +impl AcpClientLoop { + fn new(config: AcpProviderConfig, goose_mode: Arc>) -> Self { + Self { + config, + goose_mode, + prompt_response_tx: Arc::new(Mutex::new(None)), } - }; + } - match run_protocol_loop_with_child(config, child, &mut rx, init_tx.clone()).await { - Ok(()) => tracing::debug!("ACP protocol loop exited cleanly"), - Err(e) => { - let message = e.to_string(); - tracing::error!(error = %e, "ACP protocol loop error"); - send_init_result(&init_tx, Err(anyhow::anyhow!(message))); + async fn spawn( + self, + mut rx: mpsc::Receiver, + init_tx: oneshot::Sender>, + ) { + let child = match spawn_acp_process(&self.config).await { + Ok(c) => c, + Err(e) => { + let _ = init_tx.send(Err(anyhow::anyhow!("{e}"))); + tracing::error!("failed to spawn ACP process: {e}"); + return; + } + }; + + match self.run_with_child(child, &mut rx, init_tx).await { + Ok(()) => tracing::debug!("ACP protocol loop exited cleanly"), + Err(e) => tracing::error!(error = %e, "ACP protocol loop error"), } } + + async fn run_with_child( + self, + mut child: Child, + rx: &mut mpsc::Receiver, + init_tx: oneshot::Sender>, + ) -> Result<()> { + let stdin = child.stdin.take().context("no stdin")?; + let stdout = child.stdout.take().context("no stdout")?; + let transport = sacp::ByteStreams::new(stdin.compat_write(), stdout.compat()); + self.run(transport, rx, init_tx).await + } + + async fn run( + self, + transport: sacp::ByteStreams, + rx: &mut mpsc::Receiver, + init_tx: oneshot::Sender>, + ) -> Result<()> + where + R: futures::AsyncRead + Unpin + Send + 'static, + W: futures::AsyncWrite + Unpin + Send + 'static, + { + let AcpClientLoop { + config, + goose_mode, + prompt_response_tx, + } = self; + + ClientToAgent::builder() + .on_receive_notification( + { + let prompt_response_tx = prompt_response_tx.clone(); + async move |notification: SessionNotification, _cx| { + // stream() reads goose_mode at call time, so it must + // reflect any prior set_mode before the next prompt. + if let SessionUpdate::CurrentModeUpdate(update) = ¬ification.update { + if let Ok(mode) = GooseMode::from_str(&update.current_mode_id.0) { + if let Ok(mut guard) = goose_mode.lock() { + *guard = mode; + } + } + } + if let Some(tx) = prompt_response_tx + .lock() + .ok() + .as_ref() + .and_then(|g| g.as_ref()) + { + match notification.update { + SessionUpdate::AgentMessageChunk(ContentChunk { + content: ContentBlock::Text(TextContent { text, .. }), + .. + }) => { + let _ = tx.try_send(AcpUpdate::Text(text)); + } + SessionUpdate::AgentThoughtChunk(ContentChunk { + content: ContentBlock::Text(TextContent { text, .. }), + .. + }) => { + let _ = tx.try_send(AcpUpdate::Thought(text)); + } + SessionUpdate::ToolCall(tool_call) => { + let _ = tx.try_send(AcpUpdate::ToolCallStart { + id: tool_call.tool_call_id.0.to_string(), + }); + } + SessionUpdate::ToolCallUpdate(update) => { + if update.fields.status.is_some() { + let _ = tx.try_send(AcpUpdate::ToolCallComplete { + id: update.tool_call_id.0.to_string(), + }); + } + } + _ => {} + } + } + Ok(()) + } + }, + sacp::on_receive_notification!(), + ) + .on_receive_request( + { + let prompt_response_tx = prompt_response_tx.clone(); + async move |request: RequestPermissionRequest, request_cx, _connection_cx| { + let (response_tx, response_rx) = oneshot::channel(); + + let handler = prompt_response_tx + .lock() + .ok() + .as_ref() + .and_then(|g| g.as_ref().cloned()); + let tx = handler.ok_or_else(sacp::Error::internal_error)?; + + if tx.is_closed() { + return Err(sacp::Error::internal_error()); + } + + tx.try_send(AcpUpdate::PermissionRequest { + request: Box::new(request), + response_tx, + }) + .map_err(|_| sacp::Error::internal_error())?; + + let response = response_rx.await.unwrap_or_else(|_| { + RequestPermissionResponse::new(RequestPermissionOutcome::Cancelled) + }); + request_cx.respond(response) + } + }, + sacp::on_receive_request!(), + ) + .connect_to(transport)? + .run_until(move |cx: JrConnectionCx| { + handle_requests(config, cx, rx, prompt_response_tx, init_tx) + }) + .await?; + + Ok(()) + } } async fn spawn_acp_process(config: &AcpProviderConfig) -> Result { @@ -468,154 +631,52 @@ async fn spawn_acp_process(config: &AcpProviderConfig) -> Result { cmd.spawn().context("failed to spawn ACP process") } -async fn run_protocol_loop_with_child( - config: AcpProviderConfig, - mut child: Child, - rx: &mut mpsc::Receiver, - init_tx: Arc>>>>, -) -> Result<()> { - let stdin = child.stdin.take().context("no stdin")?; - let stdout = child.stdout.take().context("no stdout")?; - let transport = sacp::ByteStreams::new(stdin.compat_write(), stdout.compat()); - run_protocol_loop_with_transport(config, transport, rx, init_tx).await -} - -async fn run_protocol_loop_with_transport( - config: AcpProviderConfig, - transport: sacp::ByteStreams, - rx: &mut mpsc::Receiver, - init_tx: Arc>>>>, -) -> Result<()> -where - R: futures::AsyncRead + Unpin + Send + 'static, - W: futures::AsyncWrite + Unpin + Send + 'static, -{ - let prompt_response_tx: Arc>>> = - Arc::new(Mutex::new(None)); - - ClientToAgent::builder() - .on_receive_notification( - { - let prompt_response_tx = prompt_response_tx.clone(); - async move |notification: SessionNotification, _cx| { - if let Some(tx) = prompt_response_tx.lock().unwrap().as_ref() { - match notification.update { - SessionUpdate::AgentMessageChunk(ContentChunk { - content: ContentBlock::Text(TextContent { text, .. }), - .. - }) => { - let _ = tx.try_send(AcpUpdate::Text(text)); - } - SessionUpdate::AgentThoughtChunk(ContentChunk { - content: ContentBlock::Text(TextContent { text, .. }), - .. - }) => { - let _ = tx.try_send(AcpUpdate::Thought(text)); - } - SessionUpdate::ToolCall(tool_call) => { - let _ = tx.try_send(AcpUpdate::ToolCallStart { - id: tool_call.tool_call_id.0.to_string(), - }); - } - SessionUpdate::ToolCallUpdate(update) => { - if update.fields.status.is_some() { - let _ = tx.try_send(AcpUpdate::ToolCallComplete { - id: update.tool_call_id.0.to_string(), - }); - } - } - _ => {} - } - } - Ok(()) - } - }, - sacp::on_receive_notification!(), - ) - .on_receive_request( - { - let prompt_response_tx = prompt_response_tx.clone(); - async move |request: RequestPermissionRequest, request_cx, _connection_cx| { - let (response_tx, response_rx) = oneshot::channel(); - - let handler = prompt_response_tx.lock().unwrap().as_ref().cloned(); - let tx = handler.ok_or_else(sacp::Error::internal_error)?; - - if tx.is_closed() { - return Err(sacp::Error::internal_error()); - } - - tx.try_send(AcpUpdate::PermissionRequest { - request: Box::new(request), - response_tx, - }) - .map_err(|_| sacp::Error::internal_error())?; - - let response = response_rx.await.unwrap_or_else(|_| { - RequestPermissionResponse::new(RequestPermissionOutcome::Cancelled) - }); - request_cx.respond(response) - } - }, - sacp::on_receive_request!(), - ) - .connect_to(transport)? - .run_until({ - let prompt_response_tx = prompt_response_tx.clone(); - move |cx: JrConnectionCx| { - handle_requests(config, cx, rx, prompt_response_tx, init_tx.clone()) - } - }) - .await?; - - Ok(()) -} - async fn handle_requests( config: AcpProviderConfig, cx: JrConnectionCx, rx: &mut mpsc::Receiver, prompt_response_tx: Arc>>>, - init_tx: Arc>>>>, + init_tx: oneshot::Sender>, ) -> Result<(), sacp::Error> { + let mut init_tx = Some(init_tx); + let init_response = cx .send_request(InitializeRequest::new(ProtocolVersion::LATEST)) .block_task() .await .map_err(|err| { let message = format!("ACP initialize failed: {err}"); - send_init_result(&init_tx, Err(anyhow::anyhow!(message.clone()))); + if let Some(tx) = init_tx.take() { + let _ = tx.send(Err(anyhow::anyhow!(message.clone()))); + } sacp::Error::internal_error().data(message) })?; let mcp_capabilities = init_response.agent_capabilities.mcp_capabilities.clone(); - send_init_result(&init_tx, Ok(init_response)); + if let Some(tx) = init_tx.take() { + let _ = tx.send(Ok(init_response)); + } while let Some(request) = rx.recv().await { match request { ClientRequest::NewSession { response_tx } => { handle_new_session_request(&config, &cx, &mcp_capabilities, response_tx).await; } - ClientRequest::SetModel { - session_id, - model_id, + ClientRequest::Untyped { + method, + params, response_tx, } => { - // sacp doesn't support session/set_model as a typed request yet - let msg = sacp::UntypedMessage::new( - "session/set_model", - serde_json::json!({ - "sessionId": session_id.0, - "modelId": model_id - }), - ) - .unwrap(); - let result = cx - .send_request(msg) - .block_task() - .await - .map(|_| ()) - .map_err(|e| anyhow::anyhow!("ACP session/set_model failed: {e}")); + // Untyped because sacp doesn't have typed client requests for + // session/set_mode and session/set_model yet. + let result = match sacp::UntypedMessage::new(&method, params) { + Ok(msg) => cx + .send_request(msg) + .block_task() + .await + .map_err(anyhow::Error::from), + Err(e) => Err(anyhow::Error::from(e)), + }; let _ = response_tx.send(result); } ClientRequest::Prompt { @@ -774,15 +835,6 @@ fn filter_supported_servers( .collect() } -fn send_init_result( - init_tx: &Arc>>>>, - result: Result, -) { - if let Some(tx) = init_tx.lock().unwrap().take() { - let _ = tx.send(result); - } -} - fn messages_to_prompt(messages: &[Message]) -> Vec { let mut content_blocks = Vec::new(); diff --git a/crates/goose/src/config/goose_mode.rs b/crates/goose/src/config/goose_mode.rs index 53ca6e18..7694e5cc 100644 --- a/crates/goose/src/config/goose_mode.rs +++ b/crates/goose/src/config/goose_mode.rs @@ -1,24 +1,33 @@ use serde::{Deserialize, Serialize}; -use strum::{Display, EnumString, IntoStaticStr, VariantNames}; +use strum::{Display, EnumMessage, EnumString, IntoStaticStr, VariantNames}; +use utoipa::ToSchema; #[derive( Copy, Clone, Debug, + Default, Eq, PartialEq, Serialize, Deserialize, Display, + EnumMessage, EnumString, IntoStaticStr, VariantNames, + ToSchema, )] #[serde(rename_all = "snake_case")] #[strum(serialize_all = "snake_case")] pub enum GooseMode { + #[default] + #[strum(message = "Automatically approve tool calls")] Auto, + #[strum(message = "Ask before every tool call")] Approve, + #[strum(message = "Ask only for sensitive tool calls")] SmartApprove, + #[strum(message = "Chat only, no tool calls")] Chat, } diff --git a/crates/goose/src/providers/base.rs b/crates/goose/src/providers/base.rs index b389a2bd..7c863b89 100644 --- a/crates/goose/src/providers/base.rs +++ b/crates/goose/src/providers/base.rs @@ -8,6 +8,7 @@ use super::canonical::{map_to_canonical_model, CanonicalModelRegistry}; use super::errors::ProviderError; use super::retry::RetryConfig; use crate::config::base::ConfigValue; +use crate::config::goose_mode::GooseMode; use crate::config::ExtensionConfig; use crate::conversation::message::{Message, MessageContent}; use crate::conversation::Conversation; @@ -710,6 +711,10 @@ pub trait Provider: Send + Sync { ) -> bool { false } + + async fn update_mode(&self, _session_id: &str, _mode: GooseMode) -> Result<(), ProviderError> { + Ok(()) + } } /// A message stream yields partial text content but complete tool calls, all within the Message object diff --git a/crates/goose/tests/agent.rs b/crates/goose/tests/agent.rs index e589b2e9..d884c37a 100644 --- a/crates/goose/tests/agent.rs +++ b/crates/goose/tests/agent.rs @@ -530,7 +530,7 @@ mod tests { session_manager.clone(), PermissionManager::instance(), None, - GooseMode::Auto, + GooseMode::default(), false, GoosePlatform::GooseCli, ); diff --git a/crates/goose/tests/providers.rs b/crates/goose/tests/providers.rs index b85bc9fe..27263f6a 100644 --- a/crates/goose/tests/providers.rs +++ b/crates/goose/tests/providers.rs @@ -24,7 +24,9 @@ use goose::providers::sagemaker_tgi::SAGEMAKER_TGI_DEFAULT_MODEL; use goose::providers::snowflake::SNOWFLAKE_DEFAULT_MODEL; use goose::providers::xai::XAI_DEFAULT_MODEL; use goose::session::{SessionManager, SessionType}; -use goose_test_support::{ExpectedSessionId, McpFixture, FAKE_CODE}; +use goose_test_support::{ + EnforceSessionId, ExpectedSessionId, IgnoreSessionId, McpFixture, FAKE_CODE, +}; use std::collections::HashMap; use std::sync::{Arc, Mutex}; use tokio_util::sync::CancellationToken; @@ -114,7 +116,7 @@ struct ProviderTestConfig { image_model: Option<&'static str>, clear_env: &'static [&'static str], skip: bool, - test_session_propagation: bool, + expected_session_id: fn() -> Arc, test_permissions: bool, test_smart_approve: bool, test_context_length_exceeded: bool, @@ -136,7 +138,7 @@ impl ProviderTestConfig { image_model: None, clear_env: &[], skip: false, - test_session_propagation: true, + expected_session_id: || Arc::new(EnforceSessionId::default()), test_permissions: true, test_smart_approve: true, test_context_length_exceeded: true, @@ -184,7 +186,7 @@ impl ProviderTestConfig { let skip = which::which(binary).is_err(); Self { skip, - test_session_propagation: false, + expected_session_id: || Arc::new(IgnoreSessionId), test_smart_approve: false, test_context_length_exceeded: false, ..Self::with_llm_provider(name, model_name, &[]) @@ -205,11 +207,7 @@ impl ProviderFixture { } let guard = env_lock::lock_env(env_vars.into_iter()); - let expected_session_id = if config.test_session_propagation { - Some(ExpectedSessionId::default()) - } else { - None - }; + let expected_session_id = (config.expected_session_id)(); let mcp = McpFixture::new(expected_session_id.clone()).await; let mcp_extension = @@ -251,9 +249,7 @@ impl ProviderFixture { ) .await?; let session_id = session.id; - if let Some(ref id) = expected_session_id { - id.set(&session_id); - } + expected_session_id.set(&session_id); agent.update_provider(provider.clone(), &session_id).await?; agent .add_extension(mcp_extension, &session_id)