feat(acp): add model selection support for session/new and session/set_model (#7112)

Signed-off-by: Adrian Cole <adrian@tetrate.io>
This commit is contained in:
Adrian Cole
2026-02-10 13:23:36 +08:00
committed by GitHub
parent 4abf91e72b
commit 49ad115fb1
12 changed files with 332 additions and 36 deletions
Generated
+1
View File
@@ -4209,6 +4209,7 @@ dependencies = [
name = "goose-acp"
version = "1.23.0"
dependencies = [
"agent-client-protocol-schema",
"anyhow",
"assert-json-diff",
"async-stream",
+1
View File
@@ -19,6 +19,7 @@ goose = { path = "../goose" }
goose-mcp = { path = "../goose-mcp" }
rmcp = { workspace = true }
sacp = "10.1.0"
agent-client-protocol-schema = { version = "0.10", features = ["unstable_session_model"] }
anyhow = { workspace = true }
tokio = { workspace = true }
tokio-util = { version = "0.7.15", features = ["compat", "rt"] }
+163 -19
View File
@@ -13,6 +13,7 @@ use goose::conversation::Conversation;
use goose::mcp_utils::ToolResult;
use goose::permission::permission_confirmation::PrincipalType;
use goose::permission::{Permission, PermissionConfirmation};
use goose::providers::base::Provider;
use goose::providers::provider_registry::ProviderConstructor;
use goose::session::session_manager::SessionType;
use goose::session::{Session, SessionManager};
@@ -21,12 +22,13 @@ use sacp::schema::{
AgentCapabilities, AuthMethod, AuthenticateRequest, AuthenticateResponse, BlobResourceContents,
CancelNotification, Content, ContentBlock, ContentChunk, EmbeddedResource,
EmbeddedResourceResource, ImageContent, InitializeRequest, InitializeResponse,
LoadSessionRequest, LoadSessionResponse, McpCapabilities, McpServer, NewSessionRequest,
NewSessionResponse, PermissionOption, PermissionOptionKind, PromptCapabilities, PromptRequest,
PromptResponse, RequestPermissionOutcome, RequestPermissionRequest, ResourceLink, SessionId,
SessionNotification, SessionUpdate, StopReason, TextContent, TextResourceContents, ToolCall,
ToolCallContent, ToolCallId, ToolCallLocation, ToolCallStatus, ToolCallUpdate,
ToolCallUpdateFields, ToolKind,
LoadSessionRequest, LoadSessionResponse, McpCapabilities, McpServer, ModelId, ModelInfo,
NewSessionRequest, NewSessionResponse, PermissionOption, PermissionOptionKind,
PromptCapabilities, PromptRequest, PromptResponse, RequestPermissionOutcome,
RequestPermissionRequest, ResourceLink, SessionId, SessionModelState, SessionNotification,
SessionUpdate, 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;
@@ -48,7 +50,7 @@ pub struct GooseAcpAgent {
agent: Arc<Agent>,
provider_factory: ProviderConstructor,
config_dir: std::path::PathBuf,
provider_initialized: tokio::sync::OnceCell<String>,
provider_initialized: tokio::sync::OnceCell<Arc<dyn Provider>>,
}
fn mcp_server_to_extension_config(mcp_server: McpServer) -> Result<ExtensionConfig, String> {
@@ -266,6 +268,22 @@ async fn add_extensions(agent: &Agent, extensions: Vec<ExtensionConfig>) {
}
}
async fn build_model_state(
provider: &dyn Provider,
current_model: &str,
) -> Result<SessionModelState, sacp::Error> {
let models = provider.fetch_recommended_models().await.map_err(|e| {
sacp::Error::internal_error().data(format!("Failed to fetch models: {}", e))
})?;
Ok(SessionModelState::new(
ModelId::new(current_model),
models
.iter()
.map(|name| ModelInfo::new(ModelId::new(&**name), &**name))
.collect(),
))
}
impl GooseAcpAgent {
pub fn permission_manager(&self) -> Arc<PermissionManager> {
Arc::clone(&self.agent.config.permission_manager)
@@ -682,7 +700,7 @@ impl GooseAcpAgent {
.map_err(|e| {
sacp::Error::internal_error().data(format!("Failed to create session: {}", e))
})?;
self.ensure_provider(&goose_session).await.map_err(|e| {
let provider = self.ensure_provider(&goose_session).await.map_err(|e| {
sacp::Error::internal_error().data(format!("Failed to set provider: {}", e))
})?;
@@ -715,25 +733,28 @@ impl GooseAcpAgent {
"Session started"
);
Ok(NewSessionResponse::new(SessionId::new(goose_session.id)))
let model_state =
build_model_state(&**provider, &provider.get_model_config().model_name).await?;
Ok(NewSessionResponse::new(SessionId::new(goose_session.id)).models(model_state))
}
// Called at most once via OnceCell; returns the model_id used.
async fn create_provider(&self, session: &Session) -> Result<String> {
async fn create_provider(&self, session: &Session) -> Result<Arc<dyn Provider>> {
let config_path = self.config_dir.join(CONFIG_YAML_NAME);
let config = Config::new(&config_path, "goose")?;
let model_id = config.get_goose_model()?;
let model_config = goose::model::ModelConfig::new(&model_id)?;
let provider = (self.provider_factory)(model_config).await?;
self.agent.update_provider(provider, &session.id).await?;
Ok(model_id)
self.agent
.update_provider(provider.clone(), &session.id)
.await?;
Ok(provider)
}
async fn ensure_provider(&self, session: &Session) -> Result<()> {
async fn ensure_provider(&self, session: &Session) -> Result<&Arc<dyn Provider>> {
self.provider_initialized
.get_or_try_init(|| self.create_provider(session))
.await?;
Ok(())
.await
}
async fn on_load_session(
@@ -750,7 +771,7 @@ impl GooseAcpAgent {
sacp::Error::invalid_params()
.data(format!("Failed to load session {}: {}", session_id, e))
})?;
self.ensure_provider(&goose_session).await.map_err(|e| {
let provider = self.ensure_provider(&goose_session).await.map_err(|e| {
sacp::Error::internal_error().data(format!("Failed to set provider: {}", e))
})?;
@@ -830,7 +851,10 @@ impl GooseAcpAgent {
"Session loaded"
);
Ok(LoadSessionResponse::new())
let model_state =
build_model_state(&**provider, &provider.get_model_config().model_name).await?;
Ok(LoadSessionResponse::new().models(model_state))
}
async fn on_prompt(
@@ -928,6 +952,28 @@ impl GooseAcpAgent {
Ok(())
}
async fn on_set_model(
&self,
session_id: &str,
model_id: &str,
) -> Result<SetSessionModelResponse, sacp::Error> {
let model_config = goose::model::ModelConfig::new(model_id).map_err(|e| {
sacp::Error::internal_error().data(format!("Invalid model config: {}", e))
})?;
let provider = (self.provider_factory)(model_config).await.map_err(|e| {
sacp::Error::internal_error().data(format!("Failed to create provider: {}", e))
})?;
self.agent
.update_provider(provider, session_id)
.await
.map_err(|e| {
sacp::Error::internal_error().data(format!("Failed to update provider: {}", e))
})?;
info!(session_id = %session_id, model_id = %model_id, "Model switched");
Ok(SetSessionModelResponse::new())
}
}
pub struct GooseAcpHandler {
@@ -997,7 +1043,30 @@ impl JrMessageHandler for GooseAcpHandler {
self.agent.on_cancel(notif).await
})
.await
.done()
// HACK: sacp doesn't support session/set_model yet, so we handle it as untyped JSON.
.otherwise({
let agent = self.agent.clone();
|message: MessageCx| async move {
match message {
MessageCx::Request(req, request_cx)
if req.method == "session/set_model" =>
{
let params: SetSessionModelRequest = serde_json::from_value(req.params)
.map_err(|e| sacp::Error::invalid_params().data(e.to_string()))?;
let resp = agent
.on_set_model(&params.session_id.0, &params.model_id.0)
.await?;
let json = serde_json::to_value(resp)
.map_err(|e| sacp::Error::internal_error().data(e.to_string()))?;
request_cx.respond(json)?;
Ok(())
}
_ => Err(sacp::Error::method_not_found()),
}
}
})
.await
.map(|()| Handled::Yes)
}
}
@@ -1189,4 +1258,79 @@ print(\"hello, world\")
) {
assert_eq!(outcome_to_confirmation(&input), expected);
}
use goose::providers::errors::ProviderError;
struct MockModelProvider {
models: Result<Vec<String>, ProviderError>,
}
#[async_trait::async_trait]
impl goose::providers::base::Provider for MockModelProvider {
fn get_name(&self) -> &str {
"mock"
}
async fn complete_with_model(
&self,
_session_id: Option<&str>,
_model_config: &goose::model::ModelConfig,
_system: &str,
_messages: &[goose::conversation::message::Message],
_tools: &[rmcp::model::Tool],
) -> Result<
(
goose::conversation::message::Message,
goose::providers::base::ProviderUsage,
),
ProviderError,
> {
unimplemented!()
}
fn get_model_config(&self) -> goose::model::ModelConfig {
goose::model::ModelConfig::new_or_fail("unused")
}
async fn fetch_recommended_models(&self) -> Result<Vec<String>, ProviderError> {
self.models.clone()
}
}
#[test_case(
"model-a", Ok(vec!["model-a".into(), "model-b".into()])
=> Ok(SessionModelState::new(
ModelId::new("model-a"),
vec![ModelInfo::new(ModelId::new("model-a"), "model-a"),
ModelInfo::new(ModelId::new("model-b"), "model-b")],
))
; "returns current and available models"
)]
#[test_case(
"model-a", Ok(vec![])
=> Ok(SessionModelState::new(ModelId::new("model-a"), vec![]))
; "empty model list"
)]
#[test_case(
"model-a", Err(ProviderError::ExecutionError("fail".into()))
=> matches Err(_)
; "fetch error propagates"
)]
#[test_case(
"switched-model", Ok(vec!["model-a".into(), "switched-model".into()])
=> Ok(SessionModelState::new(
ModelId::new("switched-model"),
vec![ModelInfo::new(ModelId::new("model-a"), "model-a"),
ModelInfo::new(ModelId::new("switched-model"), "switched-model")],
))
; "current model reflects switched model"
)]
#[tokio::test]
async fn test_build_model_state(
current_model: &str,
models: Result<Vec<String>, ProviderError>,
) -> Result<SessionModelState, sacp::Error> {
let provider = MockModelProvider { models };
build_model_state(&provider, current_model).await
}
}
+103 -3
View File
@@ -8,8 +8,10 @@ use fixtures::{OpenAiFixture, PermissionDecision, Session, TestSessionConfig};
use fs_err as fs;
use goose::config::base::CONFIG_YAML_NAME;
use goose::config::GooseMode;
use goose_test_support::{ExpectedSessionId, McpFixture, FAKE_CODE};
use sacp::schema::{McpServer, McpServerHttp, ToolCallStatus};
use goose_test_support::{ExpectedSessionId, McpFixture, FAKE_CODE, TEST_MODEL};
use sacp::schema::{
McpServer, McpServerHttp, ModelId, ModelInfo, SessionModelState, ToolCallStatus,
};
pub async fn run_config_mcp<S: Session>() {
let temp_dir = tempfile::tempdir().unwrap();
@@ -18,7 +20,7 @@ pub async fn run_config_mcp<S: Session>() {
let mcp = McpFixture::new(Some(expected_session_id.clone())).await;
let config_yaml = format!(
"GOOSE_MODEL: gpt-5-nano\nextensions:\n mcp-fixture:\n enabled: true\n type: streamable_http\n name: mcp-fixture\n description: MCP fixture\n uri: \"{}\"\n",
"GOOSE_MODEL: {TEST_MODEL}\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();
@@ -255,3 +257,101 @@ pub async fn run_prompt_mcp<S: Session>() {
assert_eq!(output.text, FAKE_CODE);
expected_session_id.assert_matches(&session.id().0);
}
pub async fn run_model_list<S: Session>() {
let expected_session_id = ExpectedSessionId::default();
let openai = OpenAiFixture::new(vec![], expected_session_id.clone()).await;
let session = S::new(TestSessionConfig::default(), openai).await;
expected_session_id.set(session.id().0.to_string());
let models = session.models().unwrap();
let expected = SessionModelState::new(
ModelId::new(TEST_MODEL),
[
"gpt-5.2",
"gpt-5.2-2025-12-11",
"gpt-5.2-chat-latest",
"gpt-5.2-codex",
"gpt-5.2-pro",
"gpt-5.2-pro-2025-12-11",
"gpt-5.1",
"gpt-5.1-2025-11-13",
"gpt-5.1-chat-latest",
"gpt-5.1-codex",
"gpt-5.1-codex-max",
"gpt-5.1-codex-mini",
"gpt-5-pro",
"gpt-5-pro-2025-10-06",
"gpt-5-codex",
"gpt-5",
"gpt-5-2025-08-07",
"gpt-5-chat-latest",
"gpt-5-mini",
"gpt-5-mini-2025-08-07",
TEST_MODEL,
"gpt-5-nano-2025-08-07",
"codex-mini-latest",
"o3",
"o3-2025-04-16",
"o4-mini",
"o4-mini-2025-04-16",
"gpt-4.1",
"gpt-4.1-2025-04-14",
"gpt-4.1-mini",
"gpt-4.1-mini-2025-04-14",
"gpt-4.1-nano",
"gpt-4.1-nano-2025-04-14",
"o1-pro",
"o1-pro-2025-03-19",
"o3-mini",
"o3-mini-2025-01-31",
"o1",
"o1-2024-12-17",
"gpt-4o",
"gpt-4o-2024-05-13",
"gpt-4o-2024-08-06",
"gpt-4o-2024-11-20",
"gpt-4o-mini",
"gpt-4o-mini-2024-07-18",
"o4-mini-deep-research",
"o4-mini-deep-research-2025-06-26",
"text-embedding-3-large",
"text-embedding-3-small",
"gpt-4",
"gpt-4-0613",
"gpt-4-turbo",
"gpt-4-turbo-2024-04-09",
"gpt-3.5-turbo",
"gpt-3.5-turbo-0125",
"gpt-3.5-turbo-1106",
"text-embedding-ada-002",
]
.iter()
.map(|id| ModelInfo::new(ModelId::new(*id), *id))
.collect(),
);
assert_eq!(*models, expected);
}
pub async fn run_set_model<S: Session>() {
let expected_session_id = ExpectedSessionId::default();
let openai = OpenAiFixture::new(
vec![(
r#""model":"o4-mini""#.into(),
include_str!("../test_data/openai_basic.txt"),
)],
expected_session_id.clone(),
)
.await;
let mut session = S::new(TestSessionConfig::default(), openai).await;
expected_session_id.set(session.id().0.to_string());
session.set_model("o4-mini").await;
let output = session
.prompt("what is 1+1", PermissionDecision::Cancel)
.await;
assert_eq!(output.text, "2");
}
+16 -3
View File
@@ -11,10 +11,10 @@ use goose::providers::openai::OpenAiProvider;
use goose::providers::provider_registry::ProviderConstructor;
use goose::session_context::SESSION_ID_HEADER;
use goose_acp::server::{serve, GooseAcpAgent};
use goose_test_support::ExpectedSessionId;
use goose_test_support::{ExpectedSessionId, TEST_MODEL};
use sacp::schema::{
McpServer, PermissionOptionKind, RequestPermissionOutcome, RequestPermissionRequest,
RequestPermissionResponse, SelectedPermissionOutcome, ToolCallStatus,
RequestPermissionResponse, SelectedPermissionOutcome, SessionModelState, ToolCallStatus,
};
use std::collections::VecDeque;
use std::future::Future;
@@ -86,6 +86,17 @@ impl OpenAiFixture {
let mock_server = MockServer::start().await;
let queue = Arc::new(Mutex::new(VecDeque::from(exchanges.clone())));
// Always return the models when asked, as there is no POST data to validate
Mock::given(method("GET"))
.and(path("/v1/models"))
.respond_with(
ResponseTemplate::new(200)
.insert_header("content-type", "application/json")
.set_body_string(include_str!("../test_data/openai_models.json")),
)
.mount(&mock_server)
.await;
Mock::given(method("POST"))
.and(path("/v1/chat/completions"))
.respond_with({
@@ -188,7 +199,7 @@ pub async fn spawn_acp_server_in_process(
// ensure_provider reads the model from config lazily, so tests need a config.yaml.
let config_path = data_root.join(goose::config::base::CONFIG_YAML_NAME);
if !config_path.exists() {
fs::write(&config_path, "GOOSE_MODEL: gpt-5-nano\n").unwrap();
fs::write(&config_path, format!("GOOSE_MODEL: {TEST_MODEL}\n")).unwrap();
}
let base_url = openai_base_url.to_string();
let provider_factory: ProviderConstructor = Arc::new(move |model_config| {
@@ -249,9 +260,11 @@ pub trait Session {
where
Self: Sized;
fn id(&self) -> &sacp::schema::SessionId;
fn models(&self) -> Option<&SessionModelState>;
fn reset_openai(&self);
fn reset_permissions(&self);
async fn prompt(&mut self, text: &str, decision: PermissionDecision) -> TestOutput;
async fn set_model(&self, model_id: &str);
}
#[allow(dead_code)]
+31 -7
View File
@@ -5,9 +5,9 @@ use super::{
use async_trait::async_trait;
use goose::config::PermissionManager;
use sacp::schema::{
ContentBlock, InitializeRequest, NewSessionRequest, PromptRequest, ProtocolVersion,
RequestPermissionRequest, SessionNotification, SessionUpdate, StopReason, TextContent,
ToolCallStatus,
ContentBlock, InitializeRequest, NewSessionRequest, NewSessionResponse, PromptRequest,
ProtocolVersion, RequestPermissionRequest, SessionModelState, SessionNotification,
SessionUpdate, StopReason, TextContent, ToolCallStatus,
};
use sacp::{ClientToAgent, JrConnectionCx};
use std::sync::{Arc, Mutex};
@@ -17,6 +17,7 @@ use tokio::sync::Notify;
pub struct ClientToAgentSession {
cx: JrConnectionCx<ClientToAgent>,
session_id: sacp::schema::SessionId,
new_session_response: NewSessionResponse,
updates: Arc<Mutex<Vec<SessionNotification>>>,
permission: Arc<Mutex<PermissionDecision>>,
notify: Arc<Notify>,
@@ -50,7 +51,7 @@ impl Session for ClientToAgentSession {
let notify = Arc::new(Notify::new());
let permission = Arc::new(Mutex::new(PermissionDecision::Cancel));
let (cx, session_id) = {
let (cx, session_id, new_session_response) = {
let updates_clone = updates.clone();
let notify_clone = notify.clone();
let permission_clone = permission.clone();
@@ -60,9 +61,12 @@ impl Session for ClientToAgentSession {
Arc::new(Mutex::new(None));
let session_id_holder: Arc<Mutex<Option<sacp::schema::SessionId>>> =
Arc::new(Mutex::new(None));
let response_holder: Arc<Mutex<Option<NewSessionResponse>>> =
Arc::new(Mutex::new(None));
let cx_holder_clone = cx_holder.clone();
let session_id_holder_clone = session_id_holder.clone();
let response_holder_clone = response_holder.clone();
let (ready_tx, ready_rx) = tokio::sync::oneshot::channel();
@@ -109,7 +113,7 @@ impl Session for ClientToAgentSession {
.unwrap();
let work_dir = tempfile::tempdir().unwrap();
let session = cx
let response = cx
.send_request(
NewSessionRequest::new(work_dir.path())
.mcp_servers(mcp_servers),
@@ -119,7 +123,8 @@ impl Session for ClientToAgentSession {
.unwrap();
*cx_holder.lock().unwrap() = Some(cx.clone());
*session_id_holder.lock().unwrap() = Some(session.session_id);
*session_id_holder.lock().unwrap() = Some(response.session_id.clone());
*response_holder_clone.lock().unwrap() = Some(response);
let _ = ready_tx.send(());
std::future::pending::<Result<(), sacp::Error>>().await
@@ -136,12 +141,14 @@ impl Session for ClientToAgentSession {
let cx = cx_holder.lock().unwrap().take().unwrap();
let session_id = session_id_holder.lock().unwrap().take().unwrap();
(cx, session_id)
let new_session_response = response_holder.lock().unwrap().take().unwrap();
(cx, session_id, new_session_response)
};
Self {
cx,
session_id,
new_session_response,
updates,
permission,
notify,
@@ -155,6 +162,10 @@ impl Session for ClientToAgentSession {
&self.session_id
}
fn models(&self) -> Option<&SessionModelState> {
self.new_session_response.models.as_ref()
}
fn reset_openai(&self) {
self._openai.reset();
}
@@ -195,6 +206,19 @@ impl Session for ClientToAgentSession {
TestOutput { text, tool_status }
}
// 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 {
+12 -2
View File
@@ -3,8 +3,8 @@ use common_tests::fixtures::initialize_agent;
use common_tests::fixtures::run_test;
use common_tests::fixtures::server::ClientToAgentSession;
use common_tests::{
run_config_mcp, run_permission_persistence, run_prompt_basic, run_prompt_codemode,
run_prompt_image, run_prompt_mcp,
run_config_mcp, run_model_list, run_permission_persistence, run_prompt_basic,
run_prompt_codemode, run_prompt_image, run_prompt_mcp, run_set_model,
};
use goose::config::GooseMode;
use goose::providers::provider_registry::ProviderConstructor;
@@ -16,6 +16,16 @@ fn test_config_mcp() {
run_test(async { run_config_mcp::<ClientToAgentSession>().await });
}
#[test]
fn test_model_list() {
run_test(async { run_model_list::<ClientToAgentSession>().await });
}
#[test]
fn test_set_model() {
run_test(async { run_set_model::<ClientToAgentSession>().await });
}
#[test]
fn test_permission_persistence() {
run_test(async { run_permission_persistence::<ClientToAgentSession>().await });
File diff suppressed because one or more lines are too long
+1 -1
View File
@@ -2,4 +2,4 @@ pub mod mcp;
pub mod session;
pub use mcp::{McpFixture, FAKE_CODE, TEST_IMAGE_B64};
pub use session::{ExpectedSessionId, TEST_SESSION_ID};
pub use session::{ExpectedSessionId, TEST_MODEL, TEST_SESSION_ID};
+1
View File
@@ -1,6 +1,7 @@
use std::sync::{Arc, Mutex};
pub const TEST_SESSION_ID: &str = "test-session-id";
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";
+1 -1
View File
@@ -2,7 +2,7 @@ use reqwest::StatusCode;
use std::time::Duration;
use thiserror::Error;
#[derive(Error, Debug, PartialEq)]
#[derive(Error, Debug, Clone, PartialEq)]
pub enum ProviderError {
#[error("Authentication error: {0}")]
Authentication(String),
+1
View File
@@ -47,6 +47,7 @@ pub const OPEN_AI_KNOWN_MODELS: &[(&str, usize)] = &[
("gpt-3.5-turbo", 16_385),
("gpt-4-turbo", 128_000),
("o4-mini", 128_000),
("gpt-5-nano", 400_000),
("gpt-5.1-codex", 400_000),
("gpt-5-codex", 400_000),
];