mirror of
https://github.com/aaif-goose/goose.git
synced 2026-07-03 14:10:03 +02:00
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:
Generated
+1
@@ -4209,6 +4209,7 @@ dependencies = [
|
||||
name = "goose-acp"
|
||||
version = "1.23.0"
|
||||
dependencies = [
|
||||
"agent-client-protocol-schema",
|
||||
"anyhow",
|
||||
"assert-json-diff",
|
||||
"async-stream",
|
||||
|
||||
@@ -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
@@ -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(¶ms.session_id.0, ¶ms.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
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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
@@ -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
@@ -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 {
|
||||
|
||||
@@ -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
@@ -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,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";
|
||||
|
||||
@@ -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),
|
||||
|
||||
@@ -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),
|
||||
];
|
||||
|
||||
Reference in New Issue
Block a user