tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

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

Browse Source 
Signed-off-by: Adrian Cole <adrian@tetrate.io>
This commit is contained in:
Adrian Cole
2026-03-13 17:30:36 +08:00
committed by GitHub
parent bc7e063244
commit 1b05c92912
18 changed files with 815 additions and 372 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Cargo.lock
Generated
+1
View File
@@ -4347,6 +4347,7 @@ dependencies = [
"schemars 1.2.1",
"serde",
"serde_json",
"strum",
"tempfile",
"test-case",
"tokio",
crates/goose-acp/Cargo.toml
+1
View File
@@ -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
crates/goose-acp/src/server.rs
+140 -43
View File
@@ -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<SessionManager>,
permission_manager: Arc<PermissionManager>,
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<SessionModeState, sacp::Error> {
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<PermissionManager> {
Arc::clone(&self.permission_manager)
@@ -332,7 +351,7 @@ impl GooseAcpAgent {
builtins: Vec<String>,
data_dir: std::path::PathBuf,
config_dir: std::path::PathBuf,
goose_mode: goose::config::GooseMode,
goose_mode: GooseMode,
disable_session_naming: bool,
) -> Result<Self> {
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<Arc<dyn Provider>> {
@@ -848,6 +870,21 @@ impl GooseAcpAgent {
Ok(provider)
}
async fn get_session_agent(
&self,
session_id: &str,
cancel_token: Option<CancellationToken>,
) -> Result<Arc<Agent>, 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<McpServer>,
@@ -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<SetSessionModeResponse, sacp::Error> {
let mode = mode_id.parse::<GooseMode>().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<EmptyResponse, sacp::Error> {
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<EmptyResponse, sacp::Error> {
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<GetToolsResponse, sacp::Error> {
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<ReadResourceResponse, sacp::Error> {
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<Arc<Agent>, 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
// - _<method>: 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"),
],
)
);
}
}
crates/goose-acp/tests/common_tests/mod.rs
+233 -59
View File
@@ -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<C: Connection>() {
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<C: Connection>() {
};
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<C: Connection>() {
);
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<C: Connection>() {
..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<C: Connection>() {
pub async fn run_fs_write_text_file_false<C: Connection>() {
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<C: Connection>() {
..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<C: Connection>() {
}
pub async fn run_fs_write_text_file_true<C: Connection>() {
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<C: Connection>() {
..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<C: Connection>() {
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<C: Connection>() {
.any(|m| &*m.id.0 == "goose-provider"));
}
pub async fn run_load_mode<C: Connection>() {
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<C: Connection>() {
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<C: Connection>() {
.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<C: Connection>() {
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<C: Connection>() {
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<C: Connection>() {
}
pub async fn run_model_set<C: Connection>() {
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<C: Connection>() {
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<C: Connection>() {
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<C: Connection>() {
};
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<C: Connection>() {
}
pub async fn run_prompt_basic<C: Connection>() {
let expected_session_id = ExpectedSessionId::default();
let expected_session_id = C::expected_session_id();
let openai = OpenAiFixture::new(
vec![(
r#"</info-msg>\nwhat is 1+1""#.into(),
@@ -347,8 +475,8 @@ pub async fn run_prompt_basic<C: Connection>() {
.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<C: Connection>() {
}
pub async fn run_prompt_codemode<C: Connection>() {
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<C: Connection>() {
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<C: Connection>() {
}
pub async fn run_prompt_image<C: Connection>() {
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<C: Connection>() {
..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<C: Connection>() {
}
pub async fn run_prompt_image_attachment<C: Connection>() {
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<C: Connection>() {
.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<C: Connection>() {
}
pub async fn run_load_session_mcp<C: Connection>() {
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<C: Connection>() {
..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<C: Connection>() {
// 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<C: Connection>() {
}
pub async fn run_prompt_mcp<C: Connection>() {
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<C: Connection>() {
..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<C: Connection>() {
assert_eq!(output.text, FAKE_CODE);
expected_session_id.assert_matches(&session.session_id().0);
}
pub async fn run_mode_set_error<C: Connection>(
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::<sacp::Error>().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
});
}
};
}
crates/goose-acp/tests/custom_requests_test.rs
+13 -12
View File
@@ -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;
crates/goose-acp/tests/fixtures/mod.rs
Vendored
+15 -7
View File
@@ -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<dyn ExpectedSessionId>,
) -> 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<S> {
pub session: S,
pub models: Option<SessionModelState>,
pub modes: Option<SessionModeState>,
}
pub struct TestConnectionConfig {
pub mcp_servers: Vec<McpServer>,
pub builtins: Vec<String>,
@@ -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<dyn ExpectedSessionId>;
async fn new(config: TestConnectionConfig, openai: OpenAiFixture) -> Self;
async fn new_session(&mut self) -> (Self::Session, Option<SessionModelState>);
async fn new_session(&mut self) -> SessionResult<Self::Session>;
async fn load_session(
&mut self,
session_id: &str,
mcp_servers: Vec<McpServer>,
) -> (Self::Session, Option<SessionModelState>);
) -> SessionResult<Self::Session>;
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)]
crates/goose-acp/tests/fixtures/provider.rs
Vendored
+49 -21
View File
@@ -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<Mutex<AcpProvider>>,
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<dyn ExpectedSessionId> {
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<SessionModelState>) {
async fn new_session(&mut self) -> SessionResult<ClientToProviderSession> {
// 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<McpServer>,
) -> (ClientToProviderSession, Option<SessionModelState>) {
) -> SessionResult<ClientToProviderSession> {
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();
}
}
crates/goose-acp/tests/fixtures/server.rs
Vendored
+46 -21
View File
@@ -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<dyn ExpectedSessionId> {
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<SessionModelState>) {
async fn new_session(&mut self) -> SessionResult<ClientToAgentSession> {
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<McpServer>,
) -> (ClientToAgentSession, Option<SessionModelState>) {
) -> SessionResult<ClientToAgentSession> {
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<Mutex<Vec<SessionNotification>>>) -> String {
crates/goose-acp/tests/provider_test.rs
+16 -2
View File
@@ -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::<ClientToProviderConnection>().await });
@@ -38,6 +40,12 @@ fn test_initialize_doesnt_hit_provider() {
run_test(async { run_initialize_doesnt_hit_provider::<ClientToProviderConnection>().await });
}
#[test]
#[ignore = "TODO: implement load_session in ACP provider"]
fn test_load_mode() {
run_test(async { run_load_mode::<ClientToProviderConnection>().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::<ClientToProviderConnection>().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::<ClientToProviderConnection>().await });
}
#[test]
fn test_model_list() {
run_test(async { run_model_list::<ClientToProviderConnection>().await });
crates/goose-acp/tests/server_test.rs
+16 -2
View File
@@ -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::<ClientToAgentConnection>().await });
@@ -34,6 +36,12 @@ fn test_initialize_doesnt_hit_provider() {
run_test(async { run_initialize_doesnt_hit_provider::<ClientToAgentConnection>().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::<ClientToAgentConnection>().await });
}
#[test]
fn test_load_model() {
run_test(async { run_load_model::<ClientToAgentConnection>().await });
@@ -44,6 +52,12 @@ fn test_load_session_mcp() {
run_test(async { run_load_session_mcp::<ClientToAgentConnection>().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::<ClientToAgentConnection>().await });
}
#[test]
fn test_model_list() {
run_test(async { run_model_list::<ClientToAgentConnection>().await });
crates/goose-test-support/src/lib.rs
+3 -1
View File
@@ -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
+12 -13
View File
@@ -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
+23 -6
View File
@@ -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) {}
}
crates/goose/src/acp/provider.rs
+223 -171
View File
@@ -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<Result<NewSessionResponse>>,
},
SetModel {
session_id: SessionId,
model_id: String,
response_tx: oneshot::Sender<Result<()>>,
Untyped {
method: String,
params: serde_json::Value,
response_tx: oneshot::Sender<Result<serde_json::Value>>,
},
Prompt {
session_id: SessionId,
@@ -77,7 +78,7 @@ enum AcpUpdate {
pub struct AcpProvider {
name: String,
model: ModelConfig,
goose_mode: GooseMode,
goose_mode: Arc<Mutex<GooseMode>>,
tx: mpsc::Sender<ClientRequest>,
permission_mapping: PermissionMapping,
rejected_tool_calls: Arc<TokioMutex<HashSet<String>>>,
@@ -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<Mutex<GooseMode>>,
tx: mpsc::Sender<ClientRequest>,
permission_mapping: PermissionMapping,
rejected_tool_calls: Arc<TokioMutex<HashSet<String>>>,
@@ -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<serde_json::Value> {
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<ClientRequest>,
init_tx: oneshot::Sender<Result<InitializeResponse>>,
) {
let init_tx = Arc::new(Mutex::new(Some(init_tx)));
goose_mode: Arc<Mutex<GooseMode>>,
prompt_response_tx: Arc<Mutex<Option<mpsc::Sender<AcpUpdate>>>>,
}
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<Mutex<GooseMode>>) -> 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<ClientRequest>,
init_tx: oneshot::Sender<Result<InitializeResponse>>,
) {
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<ClientRequest>,
init_tx: oneshot::Sender<Result<InitializeResponse>>,
) -> 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<R, W>(
self,
transport: sacp::ByteStreams<W, R>,
rx: &mut mpsc::Receiver<ClientRequest>,
init_tx: oneshot::Sender<Result<InitializeResponse>>,
) -> 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) = &notification.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<ClientToAgent>| {
handle_requests(config, cx, rx, prompt_response_tx, init_tx)
})
.await?;
Ok(())
}
}
async fn spawn_acp_process(config: &AcpProviderConfig) -> Result<Child> {
@@ -468,154 +631,52 @@ async fn spawn_acp_process(config: &AcpProviderConfig) -> Result<Child> {
cmd.spawn().context("failed to spawn ACP process")
}
async fn run_protocol_loop_with_child(
config: AcpProviderConfig,
mut child: Child,
rx: &mut mpsc::Receiver<ClientRequest>,
init_tx: Arc<Mutex<Option<oneshot::Sender<Result<InitializeResponse>>>>>,
) -> 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<R, W>(
config: AcpProviderConfig,
transport: sacp::ByteStreams<W, R>,
rx: &mut mpsc::Receiver<ClientRequest>,
init_tx: Arc<Mutex<Option<oneshot::Sender<Result<InitializeResponse>>>>>,
) -> Result<()>
where
R: futures::AsyncRead + Unpin + Send + 'static,
W: futures::AsyncWrite + Unpin + Send + 'static,
{
let prompt_response_tx: Arc<Mutex<Option<mpsc::Sender<AcpUpdate>>>> =
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<ClientToAgent>| {
handle_requests(config, cx, rx, prompt_response_tx, init_tx.clone())
}
})
.await?;
Ok(())
}
async fn handle_requests(
config: AcpProviderConfig,
cx: JrConnectionCx<ClientToAgent>,
rx: &mut mpsc::Receiver<ClientRequest>,
prompt_response_tx: Arc<Mutex<Option<mpsc::Sender<AcpUpdate>>>>,
init_tx: Arc<Mutex<Option<oneshot::Sender<Result<InitializeResponse>>>>>,
init_tx: oneshot::Sender<Result<InitializeResponse>>,
) -> 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<Mutex<Option<oneshot::Sender<Result<InitializeResponse>>>>>,
result: Result<InitializeResponse>,
) {
if let Some(tx) = init_tx.lock().unwrap().take() {
let _ = tx.send(result);
}
}
fn messages_to_prompt(messages: &[Message]) -> Vec<ContentBlock> {
let mut content_blocks = Vec::new();
crates/goose/src/config/goose_mode.rs
+10 -1
View File
@@ -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,
}
crates/goose/src/providers/base.rs
+5
View File
@@ -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
crates/goose/tests/agent.rs
+1 -1
View File
@@ -530,7 +530,7 @@ mod tests {
session_manager.clone(),
PermissionManager::instance(),
None,
GooseMode::Auto,
GooseMode::default(),
false,
GoosePlatform::GooseCli,
);
crates/goose/tests/providers.rs
+8 -12
View File
@@ -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<dyn ExpectedSessionId>,
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)