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

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

Author: codefromthecrypt

Cargo.lock

@@ -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"] }

crates/goose-acp/Cargo.toml

@@ -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
+    }
 }